Repository: google/jax

Issue:
Batching rule for 'gather' not implemented
Here is a minimal example to reproduce the error.
```
import jax
import jax.numpy as np

a = np.ones((3, 3))
def f(a):
    inds = np.arange(a.shape[0])
    return a[inds]
jax.vmap(f)(a)
```
`NotImplementedError: Batching rule for 'index_take' not implemented`

So if you want to use `vmap`, the only form of indexing you can do is scalar indexing. Am I interpreting this correctly? For me that's a huge limitation.

I would try to implement the batch rule but I think it will be complicated and I have a workaround:
`np.stack([a[ind] for ind in inds])` instead of `a[inds]`. With an ideal JIT compiler those two would be equivalent, but my guess is that the performance of this workaround is not optimal.


Discussion:
Thanks for the request! Fortunately batching rules are pretty easy to write :) We just wait for people to ask for them because we love closing issues.

I started a batching rule in 87673c7 and pushed it to a branch. It should cover examples like the one you wrote. I still need to add support for (1) batching along the index arrays themselves and (2) batching for the transpose operation `index_untake`, which might be needed for vmapping gradients of functions that involve `index_take`. I think both should be pretty straightforward, so I'll get back to them soon. Oh and (3) tests!
Thank you for being so responsive! Unfortunately when I run my original example I am hitting the NotImplementedError that you added.

```
import jax
import jax.numpy as np


def index_take_batch_rule(batched_args, batch_dims, **kwargs):
    src = batched_args[0]
    src_bdim = batch_dims[0]
    idxs = batched_args[1:]
    idxs_bdims = batch_dims[1:]
    if any(bdim is not None for bdim in idxs_bdims):
        raise NotImplementedError  # TODO(mattjj)
    axes = tuple(i+1 if i >= src_bdim else i for i in kwargs['axes'])
    return jax.lax.index_take(src, idxs, axes), src_bdim

jax.lax.batching.primitive_batchers[jax.lax.index_take_p] = index_take_batch_rule

a = np.ones((3, 3))
def f(a):
    inds = np.arange(a.shape[0])
    return a[inds]
jax.vmap(f)(a)
```
`NotImplementedError`


```
$ pip freeze | grep jax
jax==0.1.16
jaxlib==0.1.4
```
This seems to work for that case:
```
def index_take_batch_rule(batched_args, batch_dims, **kwargs):
    src = batched_args[0]
    src_bdim = batch_dims[0]
    idxs = batched_args[1:]
    idxs_bdims = batch_dims[1:]
    axes = tuple(i for i in kwargs['axes'])
    return jax.lax.index_take(src, idxs, axes), src_bdim
```
Hm, the example works for me on the branch, and in this code the index array shouldn't have a batch dim on it. Weird! Maybe there's a whitespace issue?

In any case, I'll dig in more when back at a keyboard, either later tonight or tomorrow.
Maybe it doesn't work for me because I'm missing some other recent commit.

I'm using a jax wheel from [here](https://storage.googleapis.com/jax-wheels) and just pasted in your changes from 87673c7 (I would build from source if I could but I can't get bazel to trust the self-signed cert in my network).

Edit: fixed my bazel issue ([link](https://groups.google.com/forum/#!topic/bazel-discuss/13uPDObyfQg) for posterity). Will see if a source build fixes this for me.
This stuff doesn't really involve any compiled code, so I don't think rebuilding jaxlib should be necessary.

I pushed another commit that fixed an issue, but it's still a WIP. (I've only got a few minutes here and there to work on it.)
Since #307 removed index_take as a primitive, running the OP's code now results in 

```
NotImplementedError: Batching rule for 'gather' not implemented
```

I changed the title of the issue to reflect the new topic: we want a batching rule for `lax.gather`!

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -2079,12 +2079,77 @@ def _gather_transpose_rule(t, operand, start_indices, dimension_numbers,
     index_vector_dim=dimension_numbers.index_vector_dim)
   return [scatter_add(zeros, start_indices, t, scatter_dnums), ad_util.zero]
 
+def _gather_batching_rule(batched_args, batch_dims, dimension_numbers,
+                          slice_sizes, operand_shape):
+  operand, start_indices = batched_args
+  operand_bdim, start_indices_bdim = batch_dims
+
+  if operand_bdim is not None and start_indices_bdim is None:
+    operand = batching.move_dim_to_front(operand, operand_bdim)
+    slice_sizes = (operand.shape[0],) + slice_sizes
+    offset_dims = (0,) + tuple(onp.add(1, dimension_numbers.offset_dims))
+    collapsed_slice_dims = tuple(onp.add(1, dimension_numbers.collapsed_slice_dims))
+    start_index_map = tuple(onp.add(1, dimension_numbers.start_index_map))
+    dnums = GatherDimensionNumbers(
+        offset_dims=offset_dims,
+        collapsed_slice_dims=collapsed_slice_dims,
+        start_index_map=start_index_map,
+        index_vector_dim=dimension_numbers.index_vector_dim)
+    return gather(operand, start_indices, dimension_numbers=dnums,
+                  slice_sizes=slice_sizes), 0
+
+  elif operand_bdim is None and start_indices_bdim is not None:
+    start_indices = batching.move_dim_to_front(start_indices, start_indices_bdim)
+    offset_dims = tuple(onp.add(1, dimension_numbers.offset_dims))
+    index_vector_dim = dimension_numbers.index_vector_dim + 1
+    dnums = GatherDimensionNumbers(
+        offset_dims=offset_dims,
+        collapsed_slice_dims=dimension_numbers.collapsed_slice_dims,
+        start_index_map=dimension_numbers.start_index_map,
+        index_vector_dim=index_vector_dim)
+    return gather(operand, start_indices, dimension_numbers=dnums,
+                  slice_sizes=slice_sizes), 0
+
+  else:
+    # get rid of scalar index case (noticing our start_indices.ndim is
+    # incremented by one compared to the original user code)
+    if dimension_numbers.index_vector_dim == start_indices.ndim - 1:
+      start_indices = reshape(start_indices, start_indices.shape + (1,))
+
+    # move our batch dimensions to the front to preserve sanity
+    operand = batching.move_dim_to_front(operand, operand_bdim)
+    start_indices = batching.move_dim_to_front(start_indices, start_indices_bdim)
+
+    # Example: user code had start_indices shape (3, 4, 5) and index_vector_dim
+    # of 2, and we have to deal with start_indices shape (7, 3, 4, 5). We
+    # transform that to an index_vector_dim of 3, and a start_indices of shape
+    # (7, 3, 4, 6) where we concatenated an iota that counts along our batch
+    # dimension to the front of the ndindex.
+    index_vector_dim = dimension_numbers.index_vector_dim + 1
+    count_shape = list(start_indices.shape)
+    count_shape[index_vector_dim] = 1
+    counts = broadcasted_iota(start_indices.dtype, tuple(count_shape), 0)
+    start_indices = concatenate([counts, start_indices], index_vector_dim)
+
+    slice_sizes = (1,) + slice_sizes
+    collapsed_slice_dims = (0,) + tuple(onp.add(1, dimension_numbers.collapsed_slice_dims))
+    offset_dims = tuple(onp.add(1, dimension_numbers.offset_dims))
+    start_index_map = (0,) + tuple(onp.add(1, dimension_numbers.start_index_map))
+
+    dnums = GatherDimensionNumbers(
+        offset_dims=offset_dims,
+        collapsed_slice_dims=collapsed_slice_dims,
+        start_index_map=start_index_map,
+        index_vector_dim=index_vector_dim)
+    return gather(operand, start_indices, dimension_numbers=dnums,
+                  slice_sizes=slice_sizes), 0
 
 gather_p = standard_primitive(
     _gather_shape_rule, _gather_dtype_rule, 'gather',
     _gather_translation_rule)
 ad.defjvp(gather_p, _gather_jvp_rule, None)
 ad.primitive_transposes[gather_p] = _gather_transpose_rule
+batching.primitive_batchers[gather_p] = _gather_batching_rule
 
 
 ScatterDimensionNumbers = collections.namedtuple(


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -539,6 +539,130 @@ def testLaxLinalgCholesky(self):
     expected = onp.linalg.cholesky(b)
     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.int32]
+      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 testGatherBatchedOperand(self, axis, shape, dtype, idxs, dnums,
+                               slice_sizes, rng, rng_idx):
+    fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    operand = rng(shape, dtype)
+    ans = vmap(fun, (axis, None))(operand, idxs)
+    expected = onp.stack([fun(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,
+          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.int32]
+      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 testGatherBatchedIndices(self, axis, shape, dtype, idxs, dnums,
+                               slice_sizes, rng, rng_idx):
+    fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    operand = rng(shape, dtype)
+    ans = vmap(fun, (None, axis))(operand, idxs)
+    expected = onp.stack([fun(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,
+          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 testGatherBatchedBoth(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)
+    operand = rng(shape, dtype)
+    assert operand.shape[op_axis] == idxs.shape[idxs_axis]
+    ans = vmap(fun, (op_axis, idxs_axis))(operand, idxs)
+    expected = onp.stack([fun(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],
+                    [2, 0]])
+    def f(a, ind):
+      return a[:, ind]
+    expected = onp.stack([f(a, ind[i, :]) for i in range(ind.shape[0])])
+    ans = vmap(f, (None, 0))(a, ind)
+    assert onp.all(ans == expected)
+
+  def testNumpyIndexing2(self):
+    a = np.arange(2 * 3 * 4).reshape((2, 3, 4))
+    def f(a):
+      inds = np.array([0, 2])
+      return a[:, inds]
+    ans = vmap(f)(a)
+    expected = onp.stack([f(a[:, i, :]) for i in range(a.shape[1])], axis=1)
+    assert onp.all(ans == expected)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1381,6 +1381,9 @@ def testIndexTake(self, shape, dtype, idxs, axes, rng):
           ((10, 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)),
+          ((10, 5), 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()]))