Repository: google/jax
Issue: take_along_axis inconsistent with numpy

Stumbled upon this while migrating some code for selecting a batch of main eigenvectors to jax:

```python
import numpy as onp
import jax.numpy as jnp

B, M, N = 2, 3, 5
v, m = onp.random.rand(B, N), onp.random.rand(B, M, N)

def foo(np=None):
  i = np.argmax(v, axis=-1)[..., None, None]
  return np.take_along_axis(m, i, axis=-1)
print(foo(np=onp).shape) # => (2, 3, 1)
print(foo(np=jnp).shape) # => (2, 1, 1)
```

Jax returns something of a different shape. I also found that I could omit one of the dummy axes on `i`, which numpy does not accept (`m` and `i` must have the same ndim and compatible shapes) but jax does.

Changed files:
jax/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 21ea104cb..a284fd02d 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2113,19 +2113,28 @@ def take(a, indices, axis=None, out=None, mode=None):
 
 @_wraps(getattr(onp, "take_along_axis", None))
 def take_along_axis(arr, indices, axis):
-  if axis is None and ndim(arr) != 1:
-    return take_along_axis(arr.ravel(), indices.ravel(), 0)
+  if axis is None:
+    if ndim(indices) != 1:
+      msg = "take_along_axis indices must be 1D if axis=None, got shape {}"
+      raise ValueError(msg.format(shape(indices)))
+    return take_along_axis(arr.ravel(), indices, 0)
+  elif ndim(arr) != ndim(indices):
+    msg = "indices and arr must have the same number of dimensions; {} vs. {}"
+    raise ValueError(msg.format(ndim(indices), ndim(arr)))
   elif ndim(arr) == 1:
     return lax.index_take(arr, (indices,), (0,))
   else:
     # TODO(mattjj): if we lower directly to lax.gather here, we might be able to
     # avoid the reshape on the output.
-    all_indices = [lax.broadcasted_iota(_dtype(indices), shape(indices), i)
+    arr_shape = list(shape(arr))
+    arr_shape[axis] = 1
+    out_shape = lax.broadcast_shapes(shape(indices), tuple(arr_shape))
+    all_indices = [lax.broadcasted_iota(_dtype(indices), out_shape, i)
                    for i in range(ndim(arr))]
-    all_indices[axis] = indices
+    all_indices[axis] = broadcast_to(indices, out_shape)
     all_indices = tuple(map(ravel, all_indices))
     out_flat = lax.index_take(arr, all_indices, tuple(range(ndim(arr))))
-    return reshape(out_flat, shape(indices))
+    return reshape(out_flat, out_shape)
 
 
 ### Indexing
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index f71bbea83..e0cc3cc0c 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -287,6 +287,9 @@ def _shapes_are_broadcast_compatible(shapes):
       return False
   return True
 
+def _shapes_are_equal_length(shapes):
+  return all(len(shape) == len(shapes[0]) for shape in shapes[1:])
+
 
 class LaxBackedNumpyTests(jtu.JaxTestCase):
   """Tests for LAX-backed Numpy implementation."""
@@ -1410,17 +1413,29 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_{}_axis={}".format(
-          jtu.format_shape_dtype_string(shape, dtype), axis),
-       "rng": rng, "shape": shape, "dtype": dtype, "axis": axis}
-      for shape in [(3,), (3, 4), (3, 4, 5)]
-      for axis in itertools.chain(range(len(shape)), [-1], [None])
+      {"testcase_name": "_{}_ishape={}_axis={}".format(
+          jtu.format_shape_dtype_string(x_shape, dtype), i_shape, axis),
+       "rng": rng, "x_shape": x_shape, "i_shape": i_shape, "dtype": dtype,
+       "axis": axis}
+      for x_shape, i_shape in filter(
+        _shapes_are_equal_length,
+        filter(_shapes_are_broadcast_compatible,
+               CombosWithReplacement(nonempty_nonscalar_array_shapes, 2)))
+      for axis in itertools.chain(range(len(x_shape)), [-1], [None])
       for dtype in default_dtypes
       for rng in [jtu.rand_default()]))
-  def testTakeAlongAxis(self, shape, dtype, axis, rng):
+  def testTakeAlongAxis(self, x_shape, i_shape, dtype, axis, rng):
+    i_shape = onp.array(i_shape)
+    if axis is None:
+      i_shape = [onp.prod(i_shape, dtype=onp.int64)]
+    else:
+      # Test the case where the size of the axis doesn't necessarily broadcast.
+      i_shape[axis] *= 3
+      i_shape = list(i_shape)
     def args_maker():
-      x = rng(shape, dtype)
-      i = onp.argsort(x, axis=axis)
+      x = rng(x_shape, dtype)
+      n = onp.prod(x_shape, dtype=onp.int32) if axis is None else x_shape[axis]
+      i = rng(i_shape, onp.int32) % n
       return x, i
 
     lnp_op = lambda x, i: lnp.take_along_axis(x, i, axis=axis)