Add Masked Scatter (#2663)

Co-authored-by: Awni Hannun <[email protected]>
Co-authored-by: Angelos Katharopoulos <[email protected]>
Co-authored-by: Angelos Katharopoulos <[email protected]>

benchmarks/python/blas/bench_gemv.py

@@ -1,6 +1,5 @@
 # Copyright © 2023 Apple Inc.
 
-import argparse
 import os
 import subprocess
 import time

benchmarks/python/masked_scatter.py

@@ -0,0 +1,212 @@
+import math
+import os
+import subprocess
+import time
+from copy import copy
+from functools import partial
+
+import matplotlib.pyplot as plt
+import mlx.core as mx
+import numpy as np
+import torch
+from matplotlib.ticker import FuncFormatter
+
+RESULTS_DIR = "./results"
+
+
+if not os.path.isdir(RESULTS_DIR):
+    os.mkdir(RESULTS_DIR)
+
+DEVICE_NAME = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
+DEVICE_NAME = DEVICE_NAME.decode("utf-8").strip("\n")
+
+TORCH_DEVICE = torch.device(
+    "mps"
+    if torch.backends.mps.is_available()
+    else ("cuda" if torch.cuda.is_available() else "cpu")
+)
+
+
+N_WARMUP = 5
+N_ITER_BENCH = 50
+N_ITER_FUNC = 20
+
+VECTOR_LENGTHS = [4096 * (2**i) for i in range(10)]
+MASK_DENSITIES = [0.01, 0.1, 0.25, 0.5]
+D_TYPES = ("float32", "float16")
+
+
+def _power_of_two_formatter(value, _position):
+    if value <= 0:
+        return ""
+    exponent = int(round(math.log2(value)))
+    if abs(value - (1 << exponent)) / value > 1e-6:
+        return f"{value:g}"
+    return f"$2^{{{exponent}}}$"
+
+
+def torch_sync():
+    if TORCH_DEVICE.type == "cuda":
+        torch.cuda.synchronize()
+    elif TORCH_DEVICE.type == "mps":
+        torch.mps.synchronize()
+
+
+def masked_scatter_mlx(self_arr, mask_arr, src_arr):
+    outs = []
+    for _ in range(N_ITER_FUNC):
+        out = copy(self_arr)
+        out[mask_arr] = src_arr
+        outs.append(out)
+    mx.eval(outs)
+    return outs
+
+
+@torch.no_grad()
+def masked_scatter_torch(self_tensor, mask_tensor, src_tensor):
+    outs = []
+    for _ in range(N_ITER_FUNC):
+        out = self_tensor.clone()
+        out.masked_scatter_(mask_tensor, src_tensor)
+        outs.append(out)
+    torch_sync()
+    return outs
+
+
+def measure(fn):
+    for _ in range(N_WARMUP):
+        fn()
+    start = time.perf_counter_ns()
+    for _ in range(N_ITER_BENCH):
+        fn()
+    end = time.perf_counter_ns()
+    return (end - start) * 1e-9
+
+
+def bytes_touched(length, true_count, item_size):
+    mask_bytes = length
+    self_bytes = length * item_size * 2  # read + write
+    src_bytes = true_count * item_size
+    return (mask_bytes + self_bytes + src_bytes) * N_ITER_FUNC * N_ITER_BENCH
+
+
+def build_case(length, density, np_dtype, torch_dtype):
+    true_count = max(1, int(round(length * density)))
+
+    rng = np.random.default_rng()
+    self_np = rng.normal(0.0, 1.0, length).astype(np_dtype)
+    mask_np = np.zeros(length, dtype=bool)
+    mask_np[:true_count] = True
+    rng.shuffle(mask_np)
+    src_np = rng.normal(0.0, 1.0, true_count).astype(np_dtype)
+
+    self_mlx = mx.array(self_np)
+    mask_mlx = mx.array(mask_np)
+    src_mlx = mx.array(src_np)
+
+    self_torch = torch.from_numpy(self_np).to(device=TORCH_DEVICE, dtype=torch_dtype)
+    mask_torch = torch.from_numpy(mask_np).to(device=TORCH_DEVICE)
+    src_torch = torch.from_numpy(src_np).to(device=TORCH_DEVICE, dtype=torch_dtype)
+
+    # Correctness check once per configuration
+    mx_out = mx.array(self_np)
+    mx_out[mask_mlx] = src_mlx
+    mx.eval(mx_out)
+    torch_out = self_torch.clone()
+    torch_out.masked_scatter_(mask_torch, src_torch)
+
+    atol = 5e-3 if np_dtype == np.float16 else 1e-5
+    if not np.allclose(np.array(mx_out), torch_out.cpu().numpy(), atol=atol):
+        raise AssertionError("masked_scatter results diverged between MLX and Torch")
+
+    return (self_mlx, mask_mlx, src_mlx, self_torch, mask_torch, src_torch, true_count)
+
+
+def bench_case(length, density, dtype):
+    np_dtype = getattr(np, dtype)
+    torch_dtype = getattr(torch, dtype)
+    (
+        self_mlx,
+        mask_mlx,
+        src_mlx,
+        self_torch,
+        mask_torch,
+        src_torch,
+        true_count,
+    ) = build_case(length, density, np_dtype, torch_dtype)
+
+    time_mlx = measure(partial(masked_scatter_mlx, self_mlx, mask_mlx, src_mlx))
+    time_torch = measure(
+        partial(masked_scatter_torch, self_torch, mask_torch, src_torch)
+    )
+
+    total_bytes = bytes_touched(length, true_count, np_dtype().itemsize)
+    bytes_per_gb = float(1024**3)
+    mlx_gbps = (total_bytes / bytes_per_gb) / time_mlx
+    torch_gbps = (total_bytes / bytes_per_gb) / time_torch
+
+    return time_mlx, time_torch, mlx_gbps, torch_gbps
+
+
+def plot_density(ax_perf, ax_speedup, density, dtype):
+    mlx_gbps = []
+    torch_gbps = []
+    mlx_times = []
+    torch_times = []
+
+    for length in VECTOR_LENGTHS:
+        t_mlx, t_torch, gbps_mlx, gbps_torch = bench_case(length, density, dtype)
+        mlx_gbps.append(gbps_mlx)
+        torch_gbps.append(gbps_torch)
+        mlx_times.append(t_mlx)
+        torch_times.append(t_torch)
+
+    ax_perf.plot(VECTOR_LENGTHS, mlx_gbps, "tab:blue", label="MLX")
+    ax_perf.plot(VECTOR_LENGTHS, torch_gbps, "tab:red", label="Torch")
+    ax_perf.set_xscale("log", base=2)
+    ax_perf.set_xticks(VECTOR_LENGTHS)
+    formatter = FuncFormatter(_power_of_two_formatter)
+    ax_perf.xaxis.set_major_formatter(formatter)
+    ax_perf.set_title(f"density={density:.2f}")
+    ax_perf.set_ylabel("GB/s")
+    ax_perf.grid(True, which="both", linestyle=":", alpha=0.4)
+    ax_perf.legend()
+
+    speedup = np.array(torch_times) / np.array(mlx_times)
+    ax_speedup.plot(VECTOR_LENGTHS, speedup, "tab:green")
+    ax_speedup.axhline(1.0, color="tab:gray", linestyle="--")
+    ax_speedup.set_xscale("log", base=2)
+    ax_speedup.set_xticks(VECTOR_LENGTHS)
+    ax_speedup.xaxis.set_major_formatter(formatter)
+    ax_speedup.set_ylabel("Speedup (Torch_t / MLX_t)")
+    ax_speedup.grid(True, which="both", linestyle=":", alpha=0.4)
+
+
+def main():
+    for dtype in D_TYPES:
+        fig, axs = plt.subplots(
+            len(MASK_DENSITIES),
+            2,
+            figsize=(10, 12),
+            layout="constrained",
+            sharex=True,
+        )
+
+        for i, density in enumerate(MASK_DENSITIES):
+            plot_density(axs[i][0], axs[i][1], density, dtype)
+            axs[i][0].set_xlabel("vector length")
+            axs[i][1].set_xlabel("vector length")
+
+        fig.suptitle(
+            f"{DEVICE_NAME.replace('Apple ', '')} ({TORCH_DEVICE.type}) | dtype={dtype}"
+        )
+        output_path = os.path.join(
+            RESULTS_DIR,
+            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.pdf",
+        )
+        fig.savefig(output_path)
+        plt.close(fig)
+
+
+if __name__ == "__main__":
+    main()

docs/src/usage/indexing.rst

@@ -70,7 +70,8 @@ Differences from NumPy
 
   * Indexing does not perform bounds checking. Indexing out of bounds is
     undefined behavior.
-  * Boolean mask based indexing is not yet supported.
+  * Boolean mask based indexing is supported for assignment only (see
+    :ref:`boolean-mask-assignment`).
 
 The reason for the lack of bounds checking is that exceptions cannot propagate
 from the GPU. Performing bounds checking for array indices before launching the
@@ -143,3 +144,51 @@ expected. For example:
 
 In the above ``dfdx`` will have the correct gradient, namely zeros at ``idx``
 and ones elsewhere.
+
+.. _boolean-mask-assignment:
+
+Boolean Mask Assignment
+-----------------------
+
+MLX supports boolean indices using NumPy syntax. A mask must already be
+a :class:`bool_` MLX :class:`array` or a NumPy ``ndarray`` with ``dtype=bool``.
+Other index types are routed through the standard scatter code.
+
+.. code-block:: shell
+
+   >>> a = mx.array([1.0, 2.0, 3.0])
+   >>> mask = mx.array([True, False, True])
+   >>> updates = mx.array([5.0, 6.0])
+   >>> a[mask] = updates
+   >>> a
+   array([5.0, 2.0, 6.0], dtype=float32)
+
+Scalar assignments broadcast to every ``True`` entry in ``mask``. For non-scalar
+assignments, ``updates`` must provide at least as many elements as there are
+``True`` entries in ``mask``.
+
+.. code-block:: shell
+
+   >>> a = mx.zeros((2, 3))
+   >>> mask = mx.array([[True, False, True],
+                        [False, False, True]])
+   >>> a[mask] = 1.0
+   >>> a
+   array([[1.0, 0.0, 1.0],
+          [0.0, 0.0, 1.0]], dtype=float32)
+
+Boolean masks follow NumPy semantics:
+
+- The mask shape must match the shape of the axes it indexes exactly. No mask
+  broadcasting occurs.
+- Any axes not covered by the mask are taken in full.
+
+.. code-block:: shell
+
+   >>> a = mx.arange(1000).reshape(10, 10, 10)
+   >>> a[mx.random.randn(10, 10) > 0.0] = 0  # valid: mask covers axes 0 and 1
+
+The mask of shape ``(10, 10)`` applies to the first two axes, so ``a[mask]``
+selects the 1-D slices ``a[i, j, :]`` where ``mask[i, j]`` is ``True``.
+Shapes such as ``(1, 10, 10)`` or ``(10, 10, 1)`` do not match the indexed
+axes and therefore raise errors.

mlx/backend/cpu/indexing.cpp

@@ -747,4 +747,108 @@ void ScatterAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
   });
 }
 
+template <typename T>
+void masked_scatter_impl(const array& mask, const array& src, array& out) {
+  ContiguousIterator mask_it(mask);
+  ContiguousIterator src_it(src);
+  ContiguousIterator out_it(out);
+
+  const bool* mask_ptr = mask.data<bool>();
+  const T* src_ptr = src.data<T>();
+  T* dst_ptr = out.data<T>();
+
+  const size_t batch_count = mask.shape(0);
+  const size_t mask_batch_size = mask.size() / batch_count;
+  const size_t src_batch_size = src.size() / batch_count;
+
+  for (uint b = 0; b < batch_count; ++b) {
+    size_t src_consumed = 0;
+    src_it.seek(b * src_batch_size);
+
+    for (size_t i = 0; i < mask_batch_size; ++i) {
+      if (mask_ptr[mask_it.loc]) {
+        if (src_consumed >= src_batch_size) {
+          throw std::runtime_error(
+              "[MaskedScatter::eval_cpu] Source does not have enough elements for mask.");
+        }
+        dst_ptr[out_it.loc] = src_ptr[src_it.loc];
+        src_it.step();
+        ++src_consumed;
+      }
+      mask_it.step();
+      out_it.step();
+    }
+  }
+}
+
+void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 3);
+
+  auto& dst = inputs[0];
+  auto& mask = inputs[1];
+  auto& src = inputs[2];
+
+  // Copy src into out (copy allocates memory for out)
+  auto ctype =
+      dst.flags().row_contiguous ? CopyType::Vector : CopyType::General;
+  copy_cpu(dst, out, ctype, stream());
+
+  if (mask.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(mask);
+  encoder.set_input_array(src);
+  encoder.set_output_array(out);
+  encoder.dispatch([mask = array::unsafe_weak_copy(mask),
+                    src = array::unsafe_weak_copy(src),
+                    out = array::unsafe_weak_copy(out)]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        masked_scatter_impl<bool>(mask, src, out);
+        break;
+      case uint8:
+        masked_scatter_impl<uint8_t>(mask, src, out);
+        break;
+      case uint16:
+        masked_scatter_impl<uint16_t>(mask, src, out);
+        break;
+      case uint32:
+        masked_scatter_impl<uint32_t>(mask, src, out);
+        break;
+      case uint64:
+        masked_scatter_impl<uint64_t>(mask, src, out);
+        break;
+      case int8:
+        masked_scatter_impl<int8_t>(mask, src, out);
+        break;
+      case int16:
+        masked_scatter_impl<int16_t>(mask, src, out);
+        break;
+      case int32:
+        masked_scatter_impl<int32_t>(mask, src, out);
+        break;
+      case int64:
+        masked_scatter_impl<int64_t>(mask, src, out);
+        break;
+      case float16:
+        masked_scatter_impl<float16_t>(mask, src, out);
+        break;
+      case float32:
+        masked_scatter_impl<float>(mask, src, out);
+        break;
+      case float64:
+        masked_scatter_impl<double>(mask, src, out);
+        break;
+      case bfloat16:
+        masked_scatter_impl<bfloat16_t>(mask, src, out);
+        break;
+      case complex64:
+        masked_scatter_impl<complex64_t>(mask, src, out);
+        break;
+    }
+  });
+}
+
 } // namespace mlx::core

mlx/backend/cuda/primitives.cpp

@@ -37,6 +37,7 @@ NO_GPU(Inverse)
 NO_GPU(Cholesky)
 NO_GPU_MULTI(Eig)
 NO_GPU_MULTI(Eigh)
+NO_GPU(MaskedScatter)
 
 namespace distributed {
 NO_GPU_MULTI(Send)

mlx/backend/metal/CMakeLists.txt

@@ -28,6 +28,7 @@ make_jit_source(binary_ops)
 make_jit_source(ternary_ops)
 make_jit_source(reduce_utils kernels/atomic.h kernels/reduction/ops.h)
 make_jit_source(indexing/scatter kernels/indexing/indexing.h)
+make_jit_source(indexing/masked_scatter)
 make_jit_source(indexing/gather kernels/indexing/indexing.h)
 make_jit_source(indexing/gather_front kernels/indexing/indexing.h)
 make_jit_source(indexing/gather_axis)

mlx/backend/metal/indexing.cpp

@@ -1,4 +1,5 @@
 // Copyright © 2023-2024 Apple Inc.
+
 #include <fmt/format.h>
 
 #include "mlx/backend/common/compiled.h"
@@ -8,7 +9,9 @@
 #include "mlx/backend/metal/jit/includes.h"
 #include "mlx/backend/metal/jit/indexing.h"
 #include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/scan.h"
 #include "mlx/backend/metal/utils.h"
+#include "mlx/dtype.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
 
@@ -641,4 +644,84 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   compute_encoder.dispatch_threads(grid_dims, group_dims);
 }
 
+void MaskedScatter::eval_gpu(const std::vector<array>& inputs, array& out) {
+  const array& dst = inputs[0];
+  const array& mask = inputs[1];
+  const array& src = inputs[2];
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  const size_t total = mask.size();
+  const CopyType ct = (total == 1)
+      ? CopyType::Scalar
+      : (dst.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+  copy_gpu(dst, out, ct, s);
+  if (total == 0) {
+    return;
+  }
+
+  array mask_flat = flatten_in_eval(mask, 1, -1, s);
+  if (mask_flat.data<void>() != mask.data<void>()) {
+    d.add_temporary(mask_flat, s.index);
+  }
+
+  if (!mask_flat.flags().row_contiguous) {
+    mask_flat = contiguous_copy_gpu(mask_flat, s);
+    d.add_temporary(mask_flat, s.index);
+  }
+
+  // Prefix (exclusive) of mask → scatter_offsets
+  array scatter_offsets(mask_flat.shape(), uint32, nullptr, {});
+  scatter_offsets.set_data(allocator::malloc(scatter_offsets.nbytes()));
+  d.add_temporary(scatter_offsets, s.index);
+
+  scan_gpu_inplace(
+      mask_flat,
+      scatter_offsets,
+      Scan::Sum,
+      /*axis=*/1,
+      /*reverse=*/false,
+      /*inclusive=*/false,
+      s);
+
+  // Kernel selection/build
+  static constexpr std::string_view kBaseName = "masked_assign";
+  const std::string dtype_tag = type_to_name(out.dtype());
+  const std::string value_type = get_type_string(out.dtype());
+  const std::string contiguous =
+      (src.flags().row_contiguous) ? "true" : "false";
+  const std::string kernel_name =
+      fmt::format("{}_{}_{}", kBaseName, dtype_tag, contiguous);
+
+  auto lib = d.get_library(kernel_name, [&]() {
+    std::string source = metal::utils();
+    source += metal::masked_scatter();
+    source += fmt::format(
+        std::string(masked_assign_kernel), kernel_name, value_type, contiguous);
+    return source;
+  });
+  auto kernel = d.get_kernel(kernel_name, lib);
+
+  // Binding
+  int bind_idx = 0;
+  const int ndim = static_cast<int>(src.ndim());
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(mask_flat, bind_idx++);
+  compute_encoder.set_input_array(scatter_offsets, bind_idx++);
+  compute_encoder.set_input_array(src, bind_idx++);
+  compute_encoder.set_output_array(out, bind_idx++);
+  compute_encoder.set_vector_bytes(src.shape(), bind_idx++);
+  compute_encoder.set_vector_bytes(src.strides(), bind_idx++);
+  compute_encoder.set_bytes(ndim, bind_idx++);
+  compute_encoder.set_bytes(src.size() / src.shape(0), bind_idx++);
+  compute_encoder.set_bytes(mask_flat.size() / mask.shape(0), bind_idx++);
+
+  // Dispatch
+  auto group_dims = get_block_dims(total, 1, 1);
+  MTL::Size grid_dims(total, 1, 1);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
 } // namespace mlx::core

mlx/backend/metal/jit/includes.h

@@ -11,6 +11,7 @@ const char* ternary_ops();
 const char* reduce_utils();
 const char* gather();
 const char* scatter();
+const char* masked_scatter();
 
 const char* arange();
 const char* unary();

mlx/backend/metal/jit/indexing.h

@@ -70,3 +70,7 @@ constexpr std::string_view scatter_kernels = R"(
       gid);
 }}
 )";
+
+constexpr std::string_view masked_assign_kernel = R"(
+template [[host_name("{0}")]] [[kernel]] decltype(masked_assign_impl<{1}, {2}>) masked_assign_impl<{1}, {2}>;
+)";

mlx/backend/metal/kernels/indexing/masked_scatter.h

@@ -0,0 +1,38 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+template <typename T, bool src_contiguous>
+[[kernel]] void masked_assign_impl(
+    const device bool* mask [[buffer(0)]],
+    const device uint* scatter_offsets [[buffer(1)]],
+    const device T* src [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    const constant int* src_shapes [[buffer(4)]],
+    const constant int64_t* src_strides [[buffer(5)]],
+    const constant int& src_ndim [[buffer(6)]],
+    const constant int64_t& src_batch_size [[buffer(7)]],
+    const constant int64_t& mask_batch_size [[buffer(8)]],
+    uint idx [[thread_position_in_grid]]) {
+  const bool mask_value = mask[idx];
+  if (!mask_value) {
+    return;
+  }
+
+  const uint src_index = scatter_offsets[idx];
+  if (src_index >= src_batch_size) {
+    return;
+  }
+
+  const uint batch_idx = idx / mask_batch_size;
+
+  if (src_contiguous) {
+    out[idx] = src[batch_idx * src_batch_size + src_index];
+  } else {
+    out[idx] = src[elem_to_loc<uint>(
+        batch_idx * src_batch_size + src_index,
+        src_shapes,
+        src_strides,
+        src_ndim)];
+  }
+}

mlx/backend/metal/kernels/scan.metal

@@ -51,6 +51,7 @@ using namespace metal;
   instantiate_strided_scan(reverse_exclusive_##name, itype, otype, op, false, true, nreads)
 
 instantiate_scan_helper(sum_bool__int32,         bool,        int32_t,     CumSum, 4)
+instantiate_scan_helper(sum_bool__uint32,        bool,        uint32_t,    CumSum, 4)
 instantiate_scan_helper(sum_uint8_uint8,         uint8_t,     uint8_t,     CumSum, 4)
 instantiate_scan_helper(sum_uint16_uint16,       uint16_t,    uint16_t,    CumSum, 4)
 instantiate_scan_helper(sum_uint32_uint32,       uint32_t,    uint32_t,    CumSum, 4)

mlx/backend/metal/scan.cpp

@@ -6,52 +6,40 @@
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/scan.h"
 #include "mlx/backend/metal/utils.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
-void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-
-  auto& s = stream();
+void scan_gpu_inplace(
+    array in,
+    array& out,
+    Scan::ReduceType reduce_type,
+    int axis,
+    bool reverse,
+    bool inclusive,
+    const Stream& s) {
   auto& d = metal::device(s.device);
 
-  bool donate = inputs[0].is_donatable();
-  auto in = inputs[0];
-  if (in.flags().contiguous && in.strides()[axis_] != 0) {
-    if (donate && in.itemsize() == out.itemsize()) {
-      out.copy_shared_buffer(in);
-    } else {
-      out.set_data(
-          allocator::malloc(in.data_size() * out.itemsize()),
-          in.data_size(),
-          in.strides(),
-          in.flags());
-    }
-  } else {
-    in = contiguous_copy_gpu(in, s);
-    out.copy_shared_buffer(in);
-  }
+  bool contiguous = in.strides()[axis] == 1;
 
-  bool contiguous = in.strides()[axis_] == 1;
-
-  std::string reduce_type;
-  switch (reduce_type_) {
+  std::string reduce_type_str;
+  switch (reduce_type) {
     case Scan::Sum:
-      reduce_type = "sum";
+      reduce_type_str = "sum";
       break;
     case Scan::Prod:
-      reduce_type = "prod";
+      reduce_type_str = "prod";
       break;
     case Scan::Max:
-      reduce_type = "max";
+      reduce_type_str = "max";
       break;
     case Scan::Min:
-      reduce_type = "min";
+      reduce_type_str = "min";
       break;
     case Scan::LogAddExp:
-      reduce_type = "logaddexp";
+      reduce_type_str = "logaddexp";
       break;
   }
 
@@ -60,23 +48,23 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
       kname,
       contiguous ? "contig_" : "strided_",
       "scan_",
-      reverse_ ? "reverse_" : "",
-      (inclusive_) ? "inclusive_" : "exclusive_",
-      reduce_type,
+      reverse ? "reverse_" : "",
+      inclusive ? "inclusive_" : "exclusive_",
+      reduce_type_str,
       "_",
       type_to_name(in),
       "_",
       type_to_name(out));
 
   auto kernel =
-      get_scan_kernel(d, kname, reverse_, inclusive_, reduce_type, in, out);
+      get_scan_kernel(d, kname, reverse, inclusive, reduce_type_str, in, out);
 
   if (contiguous) {
     auto& compute_encoder = d.get_command_encoder(s.index);
     compute_encoder.set_compute_pipeline_state(kernel);
     compute_encoder.set_input_array(in, 0);
     compute_encoder.set_output_array(out, 1);
-    size_t size = in.shape(axis_);
+    size_t size = in.shape(axis);
     compute_encoder.set_bytes(size, 2);
 
     // Compute the thread grid
@@ -95,7 +83,7 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
         thread_group_size,
         static_cast<int>(kernel->maxTotalThreadsPerThreadgroup()));
     auto tmp_grid_dims =
-        get_2d_grid_dims(in.shape(), in.strides(), /** divisor= */ size);
+        get_2d_grid_dims(in.shape(), in.strides(), /*divisor=*/size);
     MTL::Size grid_dims(
         thread_group_size, tmp_grid_dims.width, tmp_grid_dims.height);
     MTL::Size group_dims(thread_group_size, 1, 1);
@@ -106,8 +94,8 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
     compute_encoder.set_input_array(
         in.data_shared_ptr() == nullptr ? out : in, 0);
     compute_encoder.set_output_array(out, 1);
-    size_t size = in.shape(axis_);
-    size_t stride = in.strides()[axis_];
+    size_t size = in.shape(axis);
+    size_t stride = in.strides()[axis];
     int bn = 32;
     size_t stride_blocks = (stride + bn - 1) / bn;
     compute_encoder.set_bytes(size, 2);
@@ -118,8 +106,8 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
     int n_reads = (in.itemsize() <= 4) ? 4 : 2;
     int n_simdgroups = bn / n_reads;
     int thread_group_size = n_simdgroups * 32;
-    auto tmp_grid_dims = get_2d_grid_dims(
-        in.shape(), in.strides(), /** divisor= */ size * stride);
+    auto tmp_grid_dims =
+        get_2d_grid_dims(in.shape(), in.strides(), /*divisor=*/size * stride);
     if (tmp_grid_dims.width * stride_blocks <= UINT_MAX) {
       tmp_grid_dims.width *= stride_blocks;
     } else {
@@ -132,4 +120,27 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
+void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  auto in = inputs[0];
+  if (in.flags().contiguous && in.strides()[axis_] != 0) {
+    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(
+          allocator::malloc(in.data_size() * out.itemsize()),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+    }
+  } else {
+    in = contiguous_copy_gpu(in, stream());
+    out.copy_shared_buffer(in);
+  }
+
+  scan_gpu_inplace(
+      in, out, reduce_type_, axis_, reverse_, inclusive_, stream());
+}
+
 } // namespace mlx::core

mlx/backend/metal/scan.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void scan_gpu_inplace(
+    array in,
+    array& out,
+    Scan::ReduceType reduce_type,
+    int axis,
+    bool reverse,
+    bool inclusive,
+    const Stream& s);
+
+} // namespace mlx::core

mlx/backend/no_cpu/primitives.cpp

@@ -87,6 +87,7 @@ NO_CPU(LogSumExp)
 NO_CPU_MULTI(LUF)
 NO_CPU(Matmul)
 NO_CPU(Maximum)
+NO_CPU(MaskedScatter)
 NO_CPU(Minimum)
 NO_CPU(Multiply)
 NO_CPU(Negative)

mlx/backend/no_gpu/primitives.cpp

@@ -154,6 +154,7 @@ NO_GPU(Cholesky)
 NO_GPU_MULTI(Eigh)
 NO_GPU_MULTI(Eig)
 NO_GPU(View)
+NO_GPU(MaskedScatter)
 
 namespace fast {
 NO_GPU_USE_FALLBACK(LayerNorm)

mlx/ops.cpp

@@ -3458,6 +3458,91 @@ array scatter_min(
   return scatter(a, indices, updates, axes, Scatter::Min, s);
 }
 
+array masked_scatter(
+    const array& a,
+    const array& mask,
+    const array& value,
+    StreamOrDevice s /* =  {} */) {
+  if (mask.dtype() != bool_) {
+    throw std::invalid_argument("[masked_scatter] The mask has to be boolean.");
+  }
+  if (mask.ndim() == 0) {
+    throw std::invalid_argument(
+        "[masked_scatter] Scalar masks are not supported.");
+  } else if (mask.ndim() > a.ndim()) {
+    throw std::invalid_argument(
+        "[masked_scatter] The mask cannot have more dimensions than the target.");
+  }
+
+  int unmasked_dims = a.ndim() - mask.ndim();
+
+  if (value.ndim() > unmasked_dims + 1) {
+    std::ostringstream msg;
+    msg << "[masked_scatter] Value array shape must be broadcastable with the last "
+        << unmasked_dims << " dimensions of the input.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Check if the start of the mask is compatible
+  if (!std::equal(
+          mask.shape().begin(), mask.shape().end(), a.shape().begin())) {
+    std::ostringstream msg;
+    msg << "[masked_scatter] The boolean mask should have the same shape as the "
+        << "beginning of the indexed array but the mask has shape "
+        << mask.shape() << " and the array has shape " << a.shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  array expanded_mask = mask;
+  array expanded_value = astype(value, a.dtype(), s);
+
+  // Broadcast both the mask with the last unmasked_dims of a
+  if (unmasked_dims > 0) {
+    auto mask_shape = mask.shape();
+    while (mask_shape.size() < a.ndim()) {
+      mask_shape.push_back(1);
+    }
+    expanded_mask = broadcast_to(reshape(mask, mask_shape, s), a.shape(), s);
+  }
+
+  // Broadcast the value with the unmasked dims plus one extra dimension of
+  // size mask.size(). If that dim is already provided leave it as is.
+  if (value.ndim() < unmasked_dims + 1) {
+    Shape value_shape(unmasked_dims + 1 - value.ndim(), 1);
+    value_shape.insert(
+        value_shape.end(), value.shape().begin(), value.shape().end());
+    expanded_value = reshape(expanded_value, value_shape, s);
+
+    value_shape[0] = mask.size();
+    for (int i = 1; i < unmasked_dims + 1; i++) {
+      value_shape[i] = a.shape(i - unmasked_dims - 1);
+    }
+    expanded_value = broadcast_to(expanded_value, value_shape, s);
+  } else if (!std::equal(
+                 value.shape().begin() + 1,
+                 value.shape().end(),
+                 a.shape().end() - unmasked_dims)) {
+    auto value_shape = value.shape();
+    for (int i = 1; i < unmasked_dims + 1; i++) {
+      value_shape[i] = a.shape(i - unmasked_dims - 1);
+    }
+    expanded_value = broadcast_to(expanded_value, value_shape, s);
+  }
+
+  array expanded_a = expand_dims(a, 0, s);
+  expanded_mask = expand_dims(expanded_mask, 0, s);
+  expanded_value = expand_dims(expanded_value, 0, s);
+
+  return squeeze(
+      array(
+          expanded_a.shape(),
+          expanded_a.dtype(),
+          std::make_shared<MaskedScatter>(to_stream(s)),
+          {expanded_a, expanded_mask, expanded_value}),
+      0,
+      s);
+}
+
 array sqrt(const array& a, StreamOrDevice s /* = {} */) {
   auto dtype = at_least_float(a.dtype());
   return array(

mlx/ops.h

@@ -1198,6 +1198,12 @@ inline array scatter_min(
   return scatter_min(a, {indices}, updates, std::vector<int>{axis}, s);
 }
 
+array masked_scatter(
+    const array& a,
+    const array& mask,
+    const array& src,
+    StreamOrDevice s = {});
+
 /** Square root the elements of an array. */
 array sqrt(const array& a, StreamOrDevice s = {});
 

mlx/primitives.cpp

@@ -1317,15 +1317,15 @@ Shape Convolution::conv_out_shape(
 
     if (pads_lo[i - 1] < 0 || pads_hi[i - 1] < 0) {
       std::ostringstream msg;
-      msg << "[conv] Padding sizes must be non-negative." << " Got padding "
+      msg << "[conv] Padding sizes must be non-negative. Got padding "
           << pads_lo << " | " << pads_hi << ".";
       throw std::invalid_argument(msg.str());
     }
 
     if (strides[i - 1] <= 0) {
       std::ostringstream msg;
-      msg << "[conv] Stride sizes must be positive." << " Got strides "
-          << strides << ".";
+      msg << "[conv] Stride sizes must be positive."
+          << " Got strides " << strides << ".";
       throw std::invalid_argument(msg.str());
     }
 
@@ -4348,6 +4348,145 @@ bool ScatterAxis::is_equivalent(const Primitive& other) const {
   return reduce_type_ == s_other.reduce_type_ && axis_ == s_other.axis_;
 }
 
+std::vector<array> MaskedScatter::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto& s = stream();
+  const array& dst = primals[0];
+  const array& mask = primals[1];
+  const array& src = primals[2];
+  const array mask_b = broadcast_to(mask, dst.shape(), s);
+  const array& cotan = cotangents[0];
+
+  std::vector<array> vjps;
+  vjps.reserve(argnums.size());
+
+  for (int arg : argnums) {
+    if (arg == 0) {
+      vjps.push_back(where(mask_b, zeros_like(cotan, s), cotan, s));
+    } else if (arg == 2) {
+      const array mask_flat = flatten(mask_b, s);
+      const array cotan_flat = flatten(cotan, s);
+
+      const array idx_src =
+          cumsum(astype(mask_flat, int32, s), 0, false, false, s);
+      const array cotan_src =
+          where(mask_flat, cotan_flat, array(0, cotan_flat.dtype()), s);
+
+      array gsrc_flat =
+          zeros({static_cast<int>(src.size())}, cotan_src.dtype(), s);
+      if (src.size() > 0) {
+        const array cotan_updates =
+            reshape(cotan_src, {static_cast<int>(idx_src.size()), 1}, s);
+        gsrc_flat = scatter_add(gsrc_flat, idx_src, cotan_updates, 0, s);
+      }
+
+      vjps.push_back(reshape(gsrc_flat, src.shape(), s));
+    } else {
+      throw std::invalid_argument(
+          "[masked_scatter] Cannot calculate VJP with respect to mask.");
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> MaskedScatter::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto& s = stream();
+  const array& dst = primals[0];
+  const array& mask = primals[1];
+  array mask_b = mask;
+  if (mask_b.ndim() < dst.ndim()) {
+    std::vector<int> axes(dst.ndim() - mask_b.ndim(), 0);
+    std::iota(axes.begin(), axes.end(), mask_b.ndim());
+    mask_b = expand_dims(mask_b, axes, s);
+  }
+
+  array out = zeros_like(dst, s);
+  for (int arg : argnums) {
+    if (arg == 0) {
+      out = where(mask_b, out, tangents[0], s);
+    } else if (arg == 2) {
+      out = array(
+          out.shape(),
+          out.dtype(),
+          std::make_shared<MaskedScatter>(s),
+          {out, mask, tangents[1]});
+    } else {
+      throw std::invalid_argument("[masked_scatter] invalid arg index in JVP");
+    }
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> MaskedScatter::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto& s = stream();
+
+  // The inputs all had batching in the 0-th dim. So vectorization amounts to
+  //  - Move the vectorized axis first
+  //  - Expand and broadcast the unvectorized inputs
+  //  - Flatten the first two dims (the new and old batch axes)
+  //  - Masked scatter
+  //  - Unflatten the vectorized axis again
+
+  // Find the batch dim if any
+  int batch_dim = -1;
+  for (int i = 0; i < axes.size(); i++) {
+    if (axes[i] >= 0) {
+      batch_dim = inputs[i].shape(axes[i]);
+    }
+  }
+
+  // Early exit if it's not vmapped
+  if (batch_dim < 0) {
+    return {
+        {array(
+            inputs[0].shape(),
+            inputs[0].dtype(),
+            std::make_shared<MaskedScatter>(to_stream(s)),
+            inputs)},
+        {-1}};
+  }
+
+  // Move vmapped axis to 0-th dim and broadcast the non-vectorized ones
+  auto v_in = inputs;
+  for (int i = 0; i < axes.size(); i++) {
+    if (axes[i] > 0) {
+      v_in[i] = moveaxis(v_in[i], axes[i], 0, s);
+    } else if (axes[i] < 0) {
+      v_in[i] = expand_dims(v_in[i], 0, s);
+      auto in_shape = v_in[i].shape();
+      in_shape[0] = batch_dim;
+      v_in[i] = broadcast_to(v_in[i], in_shape, s);
+    }
+  }
+
+  // Flatten the first 2 dims
+  for (int i = 0; i < 3; i++) {
+    v_in[i] = flatten(v_in[i], 0, 1, s);
+  }
+
+  // Masked scatter
+  const auto result_shape = v_in[0].shape();
+  const auto result_dtype = v_in[0].dtype();
+  array result(
+      result_shape,
+      result_dtype,
+      std::make_shared<MaskedScatter>(to_stream(s)),
+      std::move(v_in));
+
+  // Now unflatten so the vectorized axis is nice and separate
+  result = unflatten(result, 0, {batch_dim, -1}, s);
+
+  return {{result}, {0}};
+}
+
 std::vector<array> Sigmoid::vjp(
     const std::vector<array>& primals,
     const std::vector<array>& cotangents,
@@ -5111,14 +5250,18 @@ std::vector<array> BlockMaskedMM::vjp(
   //    - dB_m = A_m.T [..., K, M] @ dC [..., M, N]
   //    - dA = dA_m * mask_b_lhs [..., MP, KP]
   //    - dB = dB_m * mask_b_rhs [..., KP, MP]
-  //    - dmask_b_lhs = dA_m [..., M, K] * A [..., M, K] // need [..., MP, KP]
-  //    - dmask_b_rhs = dB_m [..., K, N] * B [..., K, N] // need [..., KP, NP]
+  //    - dmask_b_lhs = dA_m [..., M, K] * A [..., M, K] // need [..., MP,
+  //    KP]
+  //    - dmask_b_rhs = dB_m [..., K, N] * B [..., K, N] // need [..., KP,
+  //    NP]
   //
   // Observations:
-  //  * If dmask_b_lhs is not needed, then dA can be calulated in one go as a
-  //    as a block_masked_mm with mask_b_lhs as the out_mask without needing to
-  //    materialize the intermediate dA_m. Similar for dB.
-  //  * If dmask_b_lhs is needed, we need to materialize dA_m directly and then
+  //  * If dmask_b_lhs is not needed, then dA can be calulated in one go as
+  //  a
+  //    as a block_masked_mm with mask_b_lhs as the out_mask without needing
+  //    to materialize the intermediate dA_m. Similar for dB.
+  //  * If dmask_b_lhs is needed, we need to materialize dA_m directly and
+  //  then
   //    point-wise multiply with A. But the output needs to be padded
 
   std::vector<array> vjps;

mlx/primitives.h

@@ -1928,6 +1928,20 @@ class ScatterAxis : public UnaryPrimitive {
   int axis_;
 };
 
+class MaskedScatter : public UnaryPrimitive {
+ public:
+  explicit MaskedScatter(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP();
+  DEFINE_GRADS();
+  DEFINE_NAME(MaskedScatter);
+  DEFINE_DEFAULT_IS_EQUIVALENT();
+  DEFINE_INPUT_OUTPUT_SHAPE();
+};
+
 class Sigmoid : public UnaryPrimitive {
  public:
   explicit Sigmoid(Stream stream) : UnaryPrimitive(stream) {}

python/src/indexing.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 #include <numeric>
+#include <optional>
 #include <sstream>
 
 #include "python/src/convert.h"
@@ -885,6 +886,22 @@ auto mlx_slice_update(
   return std::make_pair(true, out);
 }
 
+std::optional<mx::array> extract_boolean_mask(const nb::object& obj) {
+  using NDArray = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>;
+  if (nb::isinstance<mx::array>(obj)) {
+    auto mask = nb::cast<mx::array>(obj);
+    if (mask.dtype() == mx::bool_) {
+      return mask;
+    }
+  } else if (nb::isinstance<NDArray>(obj)) {
+    auto mask = nb::cast<NDArray>(obj);
+    if (mask.dtype() == nb::dtype<bool>()) {
+      return nd_array_to_mlx(mask, mx::bool_);
+    }
+  }
+  return std::nullopt;
+}
+
 void mlx_set_item(
     mx::array& src,
     const nb::object& obj,
@@ -895,6 +912,13 @@ void mlx_set_item(
     return;
   }
 
+  if (auto mask = extract_boolean_mask(obj)) {
+    auto updates = to_array(v, src.dtype());
+    auto result = masked_scatter(src, *mask, updates);
+    src.overwrite_descriptor(result);
+    return;
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     auto out = scatter(src, indices, updates, axes);

python/tests/cuda_skip.py

@@ -56,4 +56,8 @@ cuda_skip = {
     "TestQuantized.test_throw",
     "TestQuantized.test_vjp_scales_biases",
     "TestExportImport.test_export_quantized_model",
+    # Masked scatter
+    "TestOps.test_masked_scatter",
+    "TestVmap.test_vmap_masked_scatter",
+    "TestArray.test_setitem_with_boolean_mask",
 }

python/tests/test_array.py

@@ -1928,6 +1928,18 @@ class TestArray(mlx_tests.MLXTestCase):
         anp[:, idx] = 4
         self.assertTrue(np.array_equal(a, anp))
 
+    def test_setitem_with_boolean_mask(self):
+        mask_np = np.zeros((10, 10), dtype=bool)
+        mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0
+
+        mask_np = np.zeros((1, 10, 10), dtype=bool)
+        with self.assertRaises(ValueError):
+            mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0
+
+        mask_np = np.zeros((10, 10, 1), dtype=bool)
+        with self.assertRaises(ValueError):
+            mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0
+
     def test_array_namespace(self):
         a = mx.array(1.0)
         api = a.__array_namespace__()

python/tests/test_ops.py

@@ -1260,7 +1260,6 @@ class TestOps(mlx_tests.MLXTestCase):
 
     def test_put_along_axis(self):
         for ax in [None, 0, 1, 2]:
-
             a_np = np.arange(16).reshape(2, 2, 4).astype(np.int32)
             a_mlx = mx.array(a_np)
 
@@ -3138,6 +3137,69 @@ class TestOps(mlx_tests.MLXTestCase):
         out = mx.depends(b, c)
         self.assertTrue(mx.array_equal(out, b))
 
+    def test_masked_scatter(self):
+        # boolean mask updates matching numpy semantics
+        a = mx.array([1.0, 2.0, 3.0])
+        mask = mx.array([True, False, True])
+        src = mx.array([5.0, 6.0])
+        expected = mx.array([5.0, 2.0, 6.0])
+        a[mask] = src
+        self.assertTrue(mx.array_equal(a, expected))
+
+        # non-boolean mask raises
+        b = mx.array([1.0, 2.0, 3.0])
+        bad_mask = mx.array([1, 0, 1])
+        src = mx.array([4.0, 5.0])
+        with self.assertRaises((TypeError, ValueError)):
+            b[bad_mask] = src
+
+        # mask matching leading dimension selects entire trailing slices
+        c = mx.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
+        mask = mx.array([True, False])
+        src = mx.array([2.0, 3.0, 4.0])
+        expected = mx.array([[2.0, 3.0, 4.0], [1.0, 1.0, 1.0]])
+        c[mask] = src
+        self.assertTrue(mx.array_equal(c, expected))
+
+        # scalar source applies to all selected entries
+        c = mx.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
+        mask = mx.array([True, False])
+        src = 2.0
+        expected = mx.array([[2.0, 2.0, 2.0], [1.0, 1.0, 1.0]])
+        c[mask] = src
+        self.assertTrue(mx.array_equal(c, expected))
+
+        # mask with no updates leaves values unchanged
+        d = mx.array([[7.0, 8.0], [9.0, 10.0]])
+        mask = mx.zeros_like(d).astype(mx.bool_)
+        src = mx.array([1.0])
+        d[mask] = src
+        self.assertTrue(mx.array_equal(d, mx.array([[7.0, 8.0], [9.0, 10.0]])))
+
+        # empty mask leaves array unchanged
+        e = mx.zeros((0,), dtype=mx.float32)
+        mask = mx.zeros((0,), dtype=mx.bool_)
+        src = mx.zeros((0,), dtype=mx.float32)
+        e[mask] = src
+        self.assertTrue(mx.array_equal(e, mx.zeros((0,), dtype=mx.float32)))
+
+        # strided target, mask, and source derived from slices
+        target = mx.arange(10.0, dtype=mx.float32)[1::2]
+        mask = mx.array(
+            [False, True, False, False, True, False, False, True, False, False],
+            dtype=mx.bool_,
+        )[1::2]
+        src = mx.arange(-4.0, 0.0, dtype=mx.float32)[::2]
+
+        target[mask] = src
+        self.assertTrue(
+            mx.array_equal(
+                target, mx.array([-4.0, 3.0, 5.0, -2.0, 9.0], dtype=mx.float32)
+            )
+        )
+
+
+class TestBroadcast(mlx_tests.MLXTestCase):
     def test_broadcast_shapes(self):
         # Basic broadcasting
         self.assertEqual(mx.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3))

python/tests/test_vmap.py

@@ -723,6 +723,93 @@ class TestVmap(mlx_tests.MLXTestCase):
         out = mx.vmap(gconv, in_axes=(0, 0))(x, w)
         self.assertTrue(mx.allclose(expected, out))
 
+    def test_vmap_masked_scatter(self):
+        def scatter_fn(x, m, src):
+            x[m] = src
+            return x
+
+        # Batched sources
+        a = mx.array([[10, 20, 30, 40], [50, 60, 70, 80]])
+        mask = mx.array([[False, True, True, True], [True, False, True, True]])
+        src = mx.array([[1, 2, 3], [4, 5, 6]])
+
+        expected = mx.array([[10, 1, 2, 3], [4, 60, 5, 6]])
+        vmap_scatter = mx.vmap(scatter_fn, in_axes=(0, 0, 0))
+        out = vmap_scatter(a, mask, src)
+        self.assertTrue(mx.array_equal(expected, out))
+
+        # Shared source across batch (matching mask populations)
+        a = mx.array([[0, 0, 0], [5, 5, 5]])
+        mask = mx.array([[True, False, True], [False, True, True]])
+        src = mx.array([9, 8])
+
+        expected = mx.array([[9, 0, 8], [5, 9, 8]])
+        vmap_scatter = mx.vmap(scatter_fn, in_axes=(0, 0, None))
+        out = vmap_scatter(a, mask, src)
+        self.assertTrue(mx.array_equal(expected, out))
+
+        # Shared destination with batched mask and sources
+        a = mx.array([10, 20, 30, 40])
+        mask = mx.array([[True, False, False, True], [False, True, True, False]])
+        src = mx.array([[1, 2], [3, 4]])
+
+        expected = mx.array([[1, 20, 30, 2], [10, 3, 4, 40]])
+        vmap_scatter = mx.vmap(scatter_fn, in_axes=(None, 0, 0))
+        out = vmap_scatter(a, mask, src)
+        self.assertTrue(mx.array_equal(expected, out))
+
+        # Shared mask across batch with batched sources
+        a = mx.array([[0, 0, 0, 0], [10, 20, 30, 40]])
+        mask = mx.array([True, False, True, False])
+        src = mx.array([[7, 8], [9, 10]])
+
+        expected = mx.array([[7, 0, 8, 0], [9, 20, 10, 40]])
+        vmap_scatter = mx.vmap(scatter_fn, in_axes=(0, None, 0))
+        out = vmap_scatter(a, mask, src)
+        self.assertTrue(mx.array_equal(expected, out))
+
+        # Uneven mask populations with scalar broadcast
+        a = mx.array([[0.0, 0.0, 0.0, 0.0], [10.0, 20.0, 30.0, 40.0]])
+        mask = mx.array([[True, False, True, True], [False, True, False, False]])
+        shared_src = mx.array(1.5)
+
+        expected = mx.array(
+            [[1.5, 0.0, 1.5, 1.5], [10.0, 1.5, 30.0, 40.0]], dtype=a.dtype
+        )
+        vmap_scatter = mx.vmap(scatter_fn, in_axes=(0, 0, None))
+        out = vmap_scatter(a, mask, shared_src)
+        self.assertTrue(mx.array_equal(expected, out))
+
+        # Shared src with identical masks must restart for each batch
+        a = mx.array([[0, 0, 0, 0, 0], [10, 20, 30, 40, 50]])
+        mask = mx.array(
+            [[True, True, True, False, False], [True, True, True, False, False]]
+        )
+        src = mx.array([1, 2, 3, 4, 5])
+
+        expected = mx.array([[1, 2, 3, 0, 0], [1, 2, 3, 40, 50]])
+        vmap_scatter = mx.vmap(scatter_fn, in_axes=(0, 0, None))
+        out = vmap_scatter(a, mask, src)
+        self.assertTrue(mx.array_equal(expected, out))
+
+        # Double vmap
+        a = mx.zeros((8, 8, 8))
+        mask = mx.random.normal((8, 8, 8)) > 0
+        src = mx.random.normal((8, 8))
+        expected = mx.stack(
+            [
+                mx.stack(
+                    [scatter_fn(a[i, j] + 0, mask[i, j], src[i]) for j in range(8)]
+                )
+                for i in range(8)
+            ]
+        )
+        double_scatter = mx.vmap(
+            mx.vmap(scatter_fn, in_axes=(0, 0, None)), in_axes=(0, 0, 0)
+        )
+        out = double_scatter(a + 0, mask, src)
+        self.assertTrue(mx.array_equal(expected, out))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

tests/autograd_tests.cpp

@@ -13,6 +13,8 @@
 #include "mlx/graph_utils.h"
 #include "mlx/mlx.h"
 
+#include "mlx/backend/cuda/cuda.h"
+
 using namespace mlx::core;
 
 TEST_CASE("test stop gradient") {
@@ -1353,3 +1355,45 @@ TEST_CASE("test grad dynamic slices") {
     CHECK(allclose(outs[1], ones({1, 2})).item<bool>());
   }
 }
+
+TEST_CASE("test masked_scatter autograd") {
+  if (cu::is_available()) {
+    INFO("Skipping masked_scatter cuda autograd tests");
+    return;
+  }
+
+  // Test jvp
+  {
+    auto self = array({10.f, 20.f, 30.f, 40.f}, {4});
+    auto mask = array({false, true, false, true}, bool_);
+    auto src = array({7.f, 8.f}, {2});
+
+    auto self_tan = array({1.f, 2.f, 3.f, 4.f}, {4});
+    auto src_tan = array({9.f, 11.f}, {2});
+
+    auto fun = [&mask](const std::vector<array>& in) {
+      return std::vector<array>{masked_scatter(in[0], mask, in[1])};
+    };
+
+    auto outs = jvp(fun, {self, src}, {self_tan, src_tan}).second;
+    CHECK_EQ(outs.size(), 1);
+    CHECK(array_equal(outs[0], array({1.f, 9.f, 3.f, 11.f}, {4})).item<bool>());
+  }
+
+  // Test vjp
+  {
+    auto self = array({10.f, 20.f, 30.f, 40.f}, {4});
+    auto mask = array({true, false, false, true}, bool_);
+    auto src = array({7.f, 8.f}, {2});
+
+    auto f_sum = [&mask](const std::vector<array>& xs) {
+      return std::vector<array>{sum(masked_scatter(xs[0], mask, xs[1]))};
+    };
+
+    auto v = vjp(f_sum, {self, src}, {array(1.f)});
+    const auto& grads = v.second;
+
+    CHECK(array_equal(grads[0], array({0.f, 1.f, 1.f, 0.f}, {4})).item<bool>());
+    CHECK(array_equal(grads[1], array({1.f, 1.f}, {2})).item<bool>());
+  }
+}

tests/ops_tests.cpp

@@ -7,6 +7,7 @@
 
 #include "doctest/doctest.h"
 
+#include "mlx/backend/cuda/cuda.h"
 #include "mlx/mlx.h"
 
 using namespace mlx::core;
@@ -2435,6 +2436,49 @@ TEST_CASE("test scatter") {
   }
 }
 
+TEST_CASE("test masked_scatter") {
+  if (cu::is_available()) {
+    INFO("Skipping masked_scatter cuda ops tests");
+    return;
+  }
+
+  // Wrong mask dtype
+  CHECK_THROWS(masked_scatter(array({1, 2}), array({1, 2}), array({1, 2})));
+
+  // Mask must be broadcastable to self array
+  CHECK_THROWS(masked_scatter(
+      array({1, 2, 3, 4}, {2, 2}),
+      array({false, true, true, false}, {4, 1}),
+      array({1, 2})));
+
+  // 1D mask
+  {
+    auto self = zeros({4}, int32);
+    auto mask = array({true, true, false, true});
+    auto source = array({1, 2, 4});
+    auto out = masked_scatter(self, mask, source);
+    CHECK(array_equal(out, array({1, 2, 0, 4})).item<bool>());
+  }
+
+  // Empty mask
+  {
+    auto self = zeros({4}, int32);
+    auto mask = array({false, false, false, false});
+    auto source = array({1, 2, 4});
+    auto out = masked_scatter(self, mask, source);
+    CHECK(array_equal(out, self).item<bool>());
+  }
+
+  // Broadcasted mask
+  {
+    auto self = zeros({2, 2}, int32);
+    auto mask = array({true, false});
+    auto source = array({5, 6, 7, 8}, {2, 2});
+    auto out = masked_scatter(self, mask, source);
+    CHECK(array_equal(out, array({5, 6, 0, 0}, {2, 2})).item<bool>());
+  }
+}
+
 TEST_CASE("test is positive infinity") {
   array x(1.0f);
   CHECK_FALSE(isposinf(x).item<bool>());