Repository: google/jax

Issue:
lax.scatter batching (vmap) rule
Follow-up from #246.


Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -2206,7 +2206,6 @@ def _scatter_jvp(primals, tangents, update_jaxpr, update_consts,
         updates_shape=updates_shape)
   return val_out, tangent_out
 
-
 def _scatter_transpose_rule(t, operand, scatter_indices, updates,
                             update_jaxpr, update_consts, dimension_numbers,
                             updates_shape):
@@ -2233,12 +2232,65 @@ def _scatter_transpose_rule(t, operand, scatter_indices, updates,
                       slice_sizes=slice_sizes)
   return [operand_t, None, update_t]
 
+def _scatter_batching_rule(batched_args, batch_dims, update_jaxpr,
+                           update_consts, dimension_numbers, updates_shape):
+  operand, scatter_indices, updates = batched_args
+  operand_bdim, scatter_indices_bdim, updates_bdim = batch_dims
+  del update_jaxpr, update_consts, updates_shape  # Unused.
+
+  # move the operand batch dim to the front if it is not None, otherwise create
+  # it at the front (so that we can scatter into it)
+  size = next(x.shape[ax] for x, ax in zip(batched_args, batch_dims)
+              if ax is not None)
+  operand = batching.bdim_at_front(operand, operand_bdim, broadcast_size=size,
+                                   force_broadcast=True)
+  operand_bdim = 0
+
+  if scatter_indices_bdim is not None and updates_bdim is None:
+    raise NotImplementedError  # TODO(mattjj,phawkins)
+  elif scatter_indices_bdim is None and updates_bdim is not None:
+    updates = batching.move_dim_to_front(updates, updates_bdim)
+    inserted_window_dims = tuple(onp.add(1, dimension_numbers.inserted_window_dims))
+    update_window_dims = (0,) + tuple(onp.add(1, dimension_numbers.update_window_dims))
+    scatter_dims_to_operand_dims = tuple(onp.add(1, dimension_numbers.scatter_dims_to_operand_dims))
+    dnums = ScatterDimensionNumbers(
+        update_window_dims=update_window_dims,
+        inserted_window_dims=inserted_window_dims,
+        scatter_dims_to_operand_dims=scatter_dims_to_operand_dims,
+        index_vector_dim=dimension_numbers.index_vector_dim)
+    return scatter_add(operand, scatter_indices, updates, dnums), 0
+  else:
+    # see the third case in _gather_batching_rule for comparison and comments
+    if dimension_numbers.index_vector_dim == scatter_indices.ndim - 1:
+      scatter_indices = reshape(scatter_indices, scatter_indices.shape + (1,))
+
+    scatter_indices = batching.move_dim_to_front(scatter_indices,
+                                                 scatter_indices_bdim)
+    updates = batching.move_dim_to_front(updates, updates_bdim)
+
+    index_vector_dim = dimension_numbers.index_vector_dim + 1
+    count_shape = list(scatter_indices.shape)
+    count_shape[index_vector_dim] = 1
+    counts = broadcasted_iota(scatter_indices.dtype, tuple(count_shape), 0)
+    scatter_indices = concatenate([counts, scatter_indices], index_vector_dim)
+
+    update_window_dims = tuple(onp.add(1, dimension_numbers.update_window_dims))
+    inserted_window_dims = (0,) + tuple(onp.add(1, dimension_numbers.inserted_window_dims))
+    scatter_dims_to_operand_dims = (0,) + tuple(onp.add(1, dimension_numbers.scatter_dims_to_operand_dims))
+
+    dnums = ScatterDimensionNumbers(
+        update_window_dims=update_window_dims,
+        inserted_window_dims=inserted_window_dims,
+        scatter_dims_to_operand_dims=scatter_dims_to_operand_dims,
+        index_vector_dim=index_vector_dim)
+    return scatter_add(operand, scatter_indices, updates, dnums), 0
 
 scatter_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-add',
     _scatter_translation_rule)
 ad.primitive_jvps[scatter_p] = _scatter_jvp
 ad.primitive_transposes[scatter_p] = _scatter_transpose_rule
+batching.primitive_batchers[scatter_p] = _scatter_batching_rule
 
 
 def _reduce_shape_rule(operand, init_value, computation, jaxpr, consts, dimensions):


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -571,6 +571,39 @@ def testGatherBatchedOperand(self, axis, shape, dtype, idxs, dnums,
                           for i in range(operand.shape[axis])])
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_shape={}_axis={}_idxs={}_dnums={}_slice_sizes={}".format(
+          jtu.format_shape_dtype_string(shape, dtype), axis, idxs, dnums,
+          slice_sizes),
+       "axis": axis, "shape": shape, "dtype": dtype, "idxs": idxs, "dnums": dnums,
+       "slice_sizes": slice_sizes, "rng": rng, "rng_idx": rng_idx}
+      for dtype in [onp.float32, onp.float64]
+      for axis, shape, idxs, dnums, slice_sizes in [
+          (0, (3, 5), onp.array([0, 2]), lax.GatherDimensionNumbers(
+            offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,),
+            index_vector_dim=1), (1,)),
+          (1, (10, 3), onp.array([0, 0, 0]), lax.GatherDimensionNumbers(
+            offset_dims=(1,), collapsed_slice_dims=(), start_index_map=(0,),
+            index_vector_dim=1), (2,)),
+          (1, (10, 3, 5), onp.array([0, 2, 1]), lax.GatherDimensionNumbers(
+            offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0,),
+            index_vector_dim=1), (1, 3)),
+          (2, (10, 5, 3), onp.array([[0, 2], [1, 0]]), lax.GatherDimensionNumbers(
+            offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0, 1),
+            index_vector_dim=1), (1, 3)),
+      ]
+      for rng_idx in [jtu.rand_int(max(shape))]
+      for rng in [jtu.rand_default()])
+  def testGatherGradBatchedOperand(self, axis, shape, dtype, idxs, dnums,
+                                   slice_sizes, rng, rng_idx):
+    fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    gfun = grad(lambda x, idx: np.sum(np.sin(fun(x, idx))))
+    operand = rng(shape, dtype)
+    ans = vmap(gfun, (axis, None))(operand, idxs)
+    expected = onp.stack([gfun(operand[(slice(None),) * axis + (i,)], idxs)
+                          for i in range(operand.shape[axis])])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   @parameterized.named_parameters(
       {"testcase_name": "_shape={}_axis={}_idxs={}_dnums={}_slice_sizes={}".format(
           jtu.format_shape_dtype_string(shape, dtype), axis, idxs, dnums,
@@ -606,6 +639,42 @@ def testGatherBatchedIndices(self, axis, shape, dtype, idxs, dnums,
                           for i in range(idxs.shape[axis])])
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_shape={}_axis={}_idxs={}_dnums={}_slice_sizes={}".format(
+          jtu.format_shape_dtype_string(shape, dtype), axis, idxs, dnums,
+          slice_sizes),
+       "axis": axis, "shape": shape, "dtype": dtype, "idxs": idxs, "dnums": dnums,
+       "slice_sizes": slice_sizes, "rng": rng, "rng_idx": rng_idx}
+      for dtype in [onp.float32, onp.float64]
+      for axis, shape, idxs, dnums, slice_sizes in [
+          (0, (5,), onp.array([[0, 2], [1, 3]]), lax.GatherDimensionNumbers(
+              offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,),
+              index_vector_dim=1), (1,)),
+          (1, (10,), onp.array([[0, 0, 0], [0, 2, 1]]).T,
+           lax.GatherDimensionNumbers(
+               offset_dims=(1,), collapsed_slice_dims=(), start_index_map=(0,),
+               index_vector_dim=1), (2,)),
+          (1, (10, 5), onp.array([[0, 2, 1], [0, 3, 3]]).T,
+           lax.GatherDimensionNumbers(
+               offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0,),
+               index_vector_dim=1), (1, 3)),
+          (0, (10, 5), onp.array([[[0, 2], [1, 0]],
+                                  [[1, 2], [0, 3]]]), lax.GatherDimensionNumbers(
+            offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0, 1),
+            index_vector_dim=1), (1, 3)),
+      ]
+      for rng_idx in [jtu.rand_int(max(shape))]
+      for rng in [jtu.rand_default()])
+  def testGatherGradBatchedIndices(self, axis, shape, dtype, idxs, dnums,
+                                   slice_sizes, rng, rng_idx):
+    fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    gfun = grad(lambda x, idx: np.sum(np.sin(fun(x, idx))))
+    operand = rng(shape, dtype)
+    ans = vmap(gfun, (None, axis))(operand, idxs)
+    expected = onp.stack([gfun(operand, idxs[(slice(None),) * axis + (i,)])
+                          for i in range(idxs.shape[axis])])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   @parameterized.named_parameters(
       {"testcase_name": "_shape={}_op_axis={}_idxs_axis={}_idxs={}_dnums={}_slice_sizes={}".format(
           jtu.format_shape_dtype_string(shape, dtype), op_axis, idxs_axis, idxs,
@@ -644,6 +713,45 @@ def testGatherBatchedBoth(self, op_axis, idxs_axis, shape, dtype, idxs, dnums,
                           for i in range(idxs.shape[idxs_axis])])
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_shape={}_op_axis={}_idxs_axis={}_idxs={}_dnums={}_slice_sizes={}".format(
+          jtu.format_shape_dtype_string(shape, dtype), op_axis, idxs_axis, idxs,
+          dnums, slice_sizes),
+       "op_axis": op_axis, "idxs_axis": idxs_axis, "shape": shape, "dtype":
+       dtype, "idxs": idxs, "dnums": dnums, "slice_sizes": slice_sizes,
+       "rng": rng, "rng_idx": rng_idx}
+      for dtype in [onp.float32, onp.int32]
+      for op_axis, idxs_axis, shape, idxs, dnums, slice_sizes in [
+          (0, 0, (2, 5), onp.array([[0, 2], [1, 3]]), lax.GatherDimensionNumbers(
+              offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,),
+              index_vector_dim=1), (1,)),
+          (1, 1, (10, 2), onp.array([[0, 0, 0], [0, 2, 1]]).T,
+           lax.GatherDimensionNumbers(
+               offset_dims=(1,), collapsed_slice_dims=(), start_index_map=(0,),
+               index_vector_dim=1), (2,)),
+          (0, 1, (2, 10, 5,), onp.array([[0, 2, 1], [0, 3, 3]]).T,
+           lax.GatherDimensionNumbers(
+               offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0,),
+               index_vector_dim=1), (1, 3)),
+          (2, 0, (10, 5, 2), onp.array([[[0, 2], [1, 0]],
+                                        [[1, 0], [2, 0]]]), lax.GatherDimensionNumbers(
+            offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0, 1),
+            index_vector_dim=1), (1, 3)),
+      ]
+      for rng_idx in [jtu.rand_int(max(shape))]
+      for rng in [jtu.rand_default()])
+  def testGatherGradBatchedBoth(self, op_axis, idxs_axis, shape, dtype, idxs, dnums,
+                            slice_sizes, rng, rng_idx):
+    fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    gfun = grad(lambda x, idx: np.sum(np.sin(fun(x, idx))))
+    operand = rng(shape, dtype)
+    assert operand.shape[op_axis] == idxs.shape[idxs_axis]
+    ans = vmap(gfun, (op_axis, idxs_axis))(operand, idxs)
+    expected = onp.stack([gfun(operand[(slice(None),) * op_axis + (i,)],
+                              idxs[(slice(None),) * idxs_axis + (i,)])
+                          for i in range(idxs.shape[idxs_axis])])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def testNumpyIndexing1(self):
     a = np.arange(2 * 3 * 4).reshape((2, 3, 4))
     ind = onp.array([[0, 1],