Repository: google/jax

Issue:
Batching rule for `dyanmic_slice` not defined
If you use indexing , e.g. `x[0]` in a function you `vmap`, it fails.


Discussion:
Related: we are also missing the XLA Gather primitive, which is the right way to write many batch dynamic slices.
+1 hawkinsp, our index_take predated the Gather HLO (and I find that HLO confusing) so it could be that index_take can be replaced entirely. I’m interested to know if you think that’s possible. 

(I would love to take this, but probably can’t for about 36 hours.)

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1752,11 +1752,37 @@ def dynamic_slice_transpose_rule(t, operand, start_indices, slice_sizes,
   zeros = broadcast(_const(t, 0), operand_shape)
   return [dynamic_update_slice(zeros, t, start_indices), ad_util.zero]
 
+def dynamic_slice_batching_rule(batched_args, batch_dims, slice_sizes,
+                                **unused_kwargs):
+  operand, start_indices = batched_args
+  op_bdim, idx_bdim = batch_dims
+
+  if idx_bdim is None:
+    new_start_indices = concatenate(
+        [start_indices[:op_bdim], _zeros(start_indices, shape=(1,)),
+         start_indices[op_bdim:]], 0)
+    new_slice_sizes = list(slice_sizes)
+    new_slice_sizes.insert(op_bdim, operand.shape[op_bdim])
+    out = dynamic_slice(operand, new_start_indices, new_slice_sizes)
+    return out, op_bdim
+  else:
+    # TODO(mattjj): add support for Gather HLO, use it here
+    start_indices = batching.bdim_at_front(start_indices, idx_bdim)
+    if op_bdim is None:
+      out = concatenate([dynamic_slice(operand, idx, slice_sizes)
+                         for idx in start_indices], 0)
+    else:
+      operand = batching.bdim_at_front(operand, op_bdim)
+      out = concatenate([dynamic_slice(op, idx, slice_sizes)
+                         for op, idx in zip(operand, start_indices)], 0)
+    return out, 0
+
 dynamic_slice_p = standard_primitive(
     dynamic_slice_shape_rule, _input_dtype, 'dynamic_slice',
     dynamic_slice_translation_rule)
 ad.defjvp(dynamic_slice_p, dynamic_slice_jvp_rule, None)
 ad.primitive_transposes[dynamic_slice_p] = dynamic_slice_transpose_rule
+batching.primitive_batchers[dynamic_slice_p] = dynamic_slice_batching_rule
 
 
 def dynamic_update_slice_shape_rule(operand, update, start_indices,


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -306,6 +306,26 @@ def fun(x, t):
                           [0., 0., 0., 0., 2.]])
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testDynamicSlice(self):
+    # test dynamic_slice via numpy indexing syntax
+    x = onp.arange(30).reshape((10, 3))
+
+    ans = vmap(lambda x, i: x[i], in_axes=(0, None))(x, 1)
+    expected = x[:, 1]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+
+    idx = onp.array([0, 1, 2, 1, 0] * 2)
+    ans = vmap(lambda x, i: x[i], in_axes=(0, 0))(x, idx)
+    expected = x[onp.arange(10), idx]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+    x = onp.arange(3)
+    idx = onp.array([0, 1, 2, 1, 0] * 2)
+    ans = vmap(lambda x, i: x[i], in_axes=(None, 0))(x, idx)
+    expected = x[idx]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()