Repository: google/jax
Pull request #21305 (2024-05-20 12:27:11)

Support scalar boolean indices in arr.at[idx].set(vals)
Part of #18353


Patch:
diff --git a/jax/_src/ops/scatter.py b/jax/_src/ops/scatter.py
index 5934ae415..667c2707d 100644
--- a/jax/_src/ops/scatter.py
+++ b/jax/_src/ops/scatter.py
@@ -103,9 +103,6 @@ def _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx,
   idx = jnp._merge_static_and_dynamic_indices(treedef, static_idx, dynamic_idx)
   indexer = jnp._index_to_gather(jnp.shape(x), idx,
                                  normalize_indices=normalize_indices)
-  # TODO(jakevdp): implement scalar boolean logic.
-  if indexer.scalar_bool_dims:
-    raise TypeError("Scalar boolean indices are not allowed in scatter.")
 
   # Avoid calling scatter if the slice shape is empty, both as a fast path and
   # to handle cases like zeros(0)[array([], int32)].
@@ -121,6 +118,9 @@ def _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx,
   if indexer.reversed_y_dims:
     y = lax.rev(y, indexer.reversed_y_dims)
 
+  if indexer.scalar_bool_dims:
+    x = lax.expand_dims(x, indexer.scalar_bool_dims)
+
   # Transpose the gather dimensions into scatter dimensions (cf.
   # lax._gather_transpose_rule)
   dnums = lax.ScatterDimensionNumbers(
@@ -133,10 +133,11 @@ def _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx,
     indices_are_sorted=indexer.indices_are_sorted or indices_are_sorted,
     unique_indices=indexer.unique_indices or unique_indices,
     mode=mode)
+  if indexer.scalar_bool_dims:
+    out = lax.squeeze(out, indexer.scalar_bool_dims)
   return lax_internal._convert_element_type(out, dtype, weak_type)
 
 
-
 def _get_identity(op, dtype):
   """Get an appropriate identity for a given operation in a given dtype."""
   if op is lax.scatter_add:


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index 03d73a4f6..d18244062 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -1053,6 +1053,34 @@ class IndexingTest(jtu.JaxTestCase):
     jnp_fun = lambda x: jnp.asarray(x)[idx]
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
 
+  @jtu.sample_product(
+      shape=[(2, 3, 4, 5)],
+      update_ndim=[0, 1, 2],
+      idx=[
+        np.index_exp[True],
+        np.index_exp[False],
+        np.index_exp[..., True],
+        np.index_exp[..., False],
+        np.index_exp[0, :2, True],
+        np.index_exp[0, :2, False],
+        np.index_exp[:2, 0, True],
+        np.index_exp[:2, 0, False],
+        np.index_exp[:2, np.array([0, 2]), True],
+        np.index_exp[np.array([1, 0]), :, True],
+        np.index_exp[True, :, True, :, np.array(True)],
+      ]
+  )
+  def testScalarBoolUpdate(self, shape, idx, update_ndim):
+    update_shape = np.zeros(shape)[idx].shape[-update_ndim:]
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng(shape, np.int32), rng(update_shape, np.int32)]
+    def np_fun(x, update):
+      x = np.array(x, copy=True)
+      x[idx] = update
+      return x
+    jnp_fun = lambda x, update: jnp.asarray(x).at[idx].set(update)
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
+
   def testFloatIndexingError(self):
     BAD_INDEX_TYPE_ERROR = "Indexer must have integer or boolean type, got indexer with type"
     with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):