[turingH]

1. The complete scatter implementation — there is no bounds check

1.1 The index helper only handles the negative case

// mlx/backend/cpu/indexing.cpp:18-26
template <typename IdxT>
inline size_t offset_neg_idx(IdxT idx, size_t size) {
  return (idx < 0) ? idx + size : idx;   // idx >= size passes through unchanged
}

template <>
inline size_t offset_neg_idx(uint32_t idx, size_t) {
  return idx;                             // uint32 indices: no bound at all
}

size is only used for the idx < 0 wrap; a positive index is never compared to
size.

1.2 The write operators

// mlx/backend/cpu/indexing.cpp:29-34
struct None {
  template <typename T>
  void operator()(T x, T* y) { (*y) = x; }   // the actual OOB WRITE
};
struct Sum {
  template <typename T>
  void operator()(T x, T* y) { (*y) += x; }  // scatter_add: read-modify-write
};

1.3 The scatter kernel — index → offset → write, with no check in between

// mlx/backend/cpu/indexing.cpp:298-336  (function signature + body)
template <typename InT, typename IdxT, typename OpT>
void scatter(const array& updates, array& out,
             const std::vector<array>& inds, const std::vector<int>& axes) {
  int nind = inds.size();
  auto inds_ndim = updates.ndim() - out.ndim();
  size_t n_updates = nind ? inds[0].size() : 1;
  Shape update_shape(updates.shape().begin() + inds_ndim, updates.shape().end());
  size_t update_size = 1;
  for (auto us : update_shape) update_size *= us;

  std::vector<ContiguousIterator> its(inds.begin(), inds.end());
  ContiguousIterator update_it(updates);
  ContiguousIterator out_it(update_shape, out.strides(), out.ndim());

  auto out_ptr = out.data<InT>();
  auto upd_ptr = updates.data<InT>();
  for (int i = 0; i < n_updates; ++i) {
    size_t out_offset = 0;
    for (int j = 0; j < inds.size(); ++j) {
      auto ax = axes[j];
      auto idx_loc = its[j].loc;
      its[j].step();
      auto idx_val = offset_neg_idx(inds[j].data<IdxT>()[idx_loc], out.shape(ax)); // (a) index, no upper bound
      out_offset += (idx_val * out.strides()[ax]);                                  // (b) offset
    }
    update_it.seek(i * update_size);
    for (int j = 0; j < update_size; ++j) {
      OpT{}(upd_ptr[update_it.loc], out_ptr + out_offset + out_it.loc);            // (c) WRITE
      update_it.step();
      out_it.step();
    }
    out_it.reset();
    update_it.reset();
  }
}

Between (a), (b) and (c) there is no idx_val >= out.shape(ax) / out_offset >= out.size() check. The kernel reaches out_ptr + out_offset directly.

1.4 eval_cpu copies src into out, then dispatches the kernel

// mlx/backend/cpu/indexing.cpp:410-440 (condensed)
void Scatter::eval_cpu(const std::vector<array>& inputs, array& out) {
  auto& src = inputs[0];
  auto& updates = inputs.back();
  copy_cpu(src, out, ..., stream());            // allocates `out` at src.shape
  ...
  encoder.dispatch([...]() mutable {
    dispatch_all_types(out.dtype(), [&](auto type_tag) {
      dispatch_scatter<T>(out, inds, updates, axes_, reduce_type_);  // -> scatter<...>()
    });
  });
}

out is sized from src.shape (16 bytes for [4] float32); the kernel then writes at
an index-derived offset with no reference back to that shape.

2. The operator code that triggers it

A developer writes an op that scatters at a computed index:

import mlx.core as mx
out = mx.zeros((4,), mx.float32)        # 4-float buffer = 16 bytes of data
out[mx.array([6], mx.int32)] = 1.0      # computed index 6 — out of bounds (>= shape 4)
mx.eval(out)                            # scatter writes out[6]

C++ (the exact reproducer that produced the trace below):

auto out = zeros({4}, float32);         // 16-byte data buffer
auto inds = array({6}, int32);          // out-of-bounds index 6
auto updates = ones({1, 1}, float32);
auto r = scatter(out, inds, updates, 0);
eval(r);                                // writes out[6]

3. The complete call stack + memory dump

out_offset = 6 * 4 = 24, i.e. 8 bytes past the end of the 16-byte out data. ASAN
flags the write in the redzone (full stack, std::__1 template frames elided):

==29790==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x6030000036c0
WRITE of size 4 at 0x6030000036c0 thread T1
    #0 mlx::core::None::operator()<float>(float, float*) indexing.cpp:31      <- the WRITE
    #1 mlx::core::scatter<float, int, mlx::core::None>(...) indexing.cpp:331  <- (c)
    #2 mlx::core::dispatch_scatter_inds<float, int>(...) indexing.cpp:349
    #3 mlx::core::dispatch_scatter<float>(...) indexing.cpp:398
    #4 mlx::core::Scatter::eval_cpu(...)::lambda(...) indexing.cpp:435
    #5 mlx::core::dispatch_all_types<...> dtype_utils.h:51
    #6 mlx::core::Scatter::eval_cpu(...) indexing.cpp:433
    #16 mlx::core::scheduler::Scheduler::enqueue(...) scheduler.cpp:105
    #23 mlx::core::scheduler::StreamThread::thread_fn() scheduler.cpp:74

0x6030000036c0 is located 8 bytes after 24-byte region [0x6030000036a0,0x6030000036b8)
allocated by thread T0 here:
    #1 mlx::core::allocator::CommonAllocator::malloc(unsigned long) allocator.cpp:102
    #9 mlx::core::set_copy_output_data(...) copy.h:38
    #10 mlx::core::copy_cpu(...) copy.cpp:241
    #11 mlx::core::Scatter::eval_cpu(...) indexing.cpp:418                      <- the copy

Shadow bytes around the buggy address:
=>0x603000003680: fd fd fa fa 00 00 00 fa[fa]fa 00 00 00 00 fa fa

The [fa] at the arrow is the poisoned redzone byte; the write at offset 24 lands on it.

4. Why the index is never checked against the buffer

1. out is a [4] float32 array → 16 bytes of logical data (Scatter::eval_cpu
   copies src into out, sizing it from src.shape).
2. scatter(out, [6], ...) reads the index 6, calls offset_neg_idx(6, 4) — since
   6 >= 0, the idx < 0 branch is not taken and 6 is returned unchanged (the
   only check the helper does is for the negative case).
3. out_offset = 6 * strides[0] = 6 * 4 = 24. No comparison against out.size() or
   out.shape(0) happens anywhere between the index read and the write.
4. None::operator()(x, out_ptr + 24) writes 4 bytes 8 bytes past the 16-byte data
   → the redzone is hit → ASAN aborts. On the default Metal backend the same write lands
   in the large mapped Metal heap and stays silent, but it is still an out-of-bounds
   write past the logical array, never rejected and never raised as IndexError.

Root cause: the scatter kernel trusts the caller's index unconditionally — there is
no idx_val >= out.shape(ax) guard anywhere in the code path from the index read to the
out_ptr + out_offset write.

5. The fix

Add the missing bounds checks to offset_neg_idx so an out-of-bounds index throws
instead of silently writing out of bounds (this single point covers scatter,
gather, take, take_along_axis, put_along_axis — the whole indexed read/write
family):

// mlx/backend/cpu/indexing.cpp:18-26 — fixed
template <typename IdxT>
inline size_t offset_neg_idx(IdxT idx, size_t size) {
  if (idx < 0) {
    if (idx < -static_cast<IdxT>(size)) {
      throw std::out_of_range("[indexing] Index out of bounds.");
    }
    return static_cast<size_t>(idx) + size;
  }
  if (static_cast<size_t>(idx) >= size) {
    throw std::out_of_range("[indexing] Index out of bounds.");
  }
  return static_cast<size_t>(idx);
}

template <>
inline size_t offset_neg_idx(uint32_t idx, size_t size) {
  if (idx >= size) {
    throw std::out_of_range("[indexing] Index out of bounds.");
  }
  return idx;
}

After the fix, the operator out[mx.array([6], mx.int32)] = 1.0 raises IndexError
(the Python mapping of std::out_of_range) at eval time instead of corrupting the heap.

6. Unit test + run commands

Regression test (fails / aborts before the fix, passes after):

TEST_CASE("test scatter oob index") {
  auto out = zeros({4}, float32);
  auto inds = array({6}, int32);        // out-of-bounds index
  auto updates = ones({1, 1}, float32);
  auto r = scatter(out, inds, updates, 0);
  CHECK_THROWS_AS(eval(r), std::out_of_range);   // before fix: heap-buffer-overflow
}

cmake -S . -B build-asan -DCMAKE_BUILD_TYPE=Debug -DMLX_BUILD_METAL=OFF \
  -DMLX_BUILD_TESTS=ON \
  -DCMAKE_C_FLAGS="-fsanitize=address -fno-omit-frame-pointer -g" \
  -DCMAKE_CXX_FLAGS="-fsanitize=address -fno-omit-frame-pointer -g" \
  -DCMAKE_EXE_LINKER_FLAGS="-fsanitize=address"
cmake --build build-asan --target tests -j8
./build-asan/tests/tests --test-case="test scatter oob index"