Merge pull request jax-ml#5768 from jakevdp:fix-take

PiperOrigin-RevId: 358049345

Author: jax authors

docs/CHANGELOG.md

@@ -17,6 +17,8 @@ PLEASE REMEMBER TO CHANGE THE '..master' WITH AN ACTUAL TAG in GITHUB LINK.
     from JAX ([#5627](https://github.com/google/jax/pull/5627)
     and [README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md#calling-tensorflow-functions-from-jax)).
 * Bug fixes:
+  * {func}`jax.numpy.take` properly handles negative indices
+    ([#5768](https://github.com/google/jax/pull/5768))
 * Breaking changes:
   * JAX's promotion rules were adjusted to make promotion more consistent and
     invariant to JIT. In particular, binary operations can now result in weakly-typed

jax/_src/numpy/lax_numpy.py

@@ -4142,12 +4142,15 @@ def take(a, indices, axis: Optional[int] = None, out=None, mode=None):
   else:
     axis_idx = _canonicalize_axis(axis, ndim(a))
 
-  if mode == "raise":
+  if mode is None:
+    # lax.gather() does not support negative indices, so we wrap them here
+    indices = where(indices < 0, indices + a.shape[axis_idx], indices)
+  elif mode == "raise":
     # TODO(phawkins): we have no way to report out of bounds errors yet.
     raise NotImplementedError("The 'raise' mode to jnp.take is not supported.")
   elif mode == "wrap":
     indices = mod(indices, _constant_like(indices, a.shape[axis_idx]))
-  elif mode != "clip" and mode is not None:
+  elif mode != "clip":
     raise ValueError("Invalid mode '{}' for np.take".format(mode))
 
   index_dims = len(shape(indices))

tests/lax_numpy_test.py

@@ -3621,15 +3621,18 @@ def testUnpackbits(self, shape, dtype, axis, bitorder, count):
                                   [cast(Optional[int], None)])
       for dtype in all_dtypes
       for index_dtype in int_dtypes
-      for mode in ['wrap', 'clip']))
+      for mode in [None, 'wrap', 'clip']))
   def testTake(self, shape, dtype, index_shape, index_dtype, axis, mode):
     def args_maker():
       x = rng(shape, dtype)
       i = rng_indices(index_shape, index_dtype)
       return x, i
 
     rng = jtu.rand_default(self.rng())
-    rng_indices = jtu.rand_int(self.rng(), -5, 5)
+    if mode is None:
+      rng_indices = jtu.rand_int(self.rng(), -shape[axis or 0], shape[axis or 0])
+    else:
+      rng_indices = jtu.rand_int(self.rng(), -5, 5)
     jnp_op = lambda x, i: jnp.take(x, i, axis=axis, mode=mode)
     np_op = lambda x, i: np.take(x, i, axis=axis, mode=mode)
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker)