Repository: google/jax

Issue:
Incorrect results from take_along_axis with uint8 indices
I get incorrect values from `jnp.take_along_axis ` for arrays larger than 128 in any dimension and `uint8 `indices. It works with when I cast the indices to larger bit widths.
Jax version 0.2.6
Jaxlib version 0.1.57

```
h = np.random.random((256,256,100))
g = np.random.randint(0,100, size=(256,256,1)).astype(np.uint8)

q0 = jnp.take_along_axis(h, g, axis=-1)
q1 = np.take_along_axis( h, g, axis=-1)

print(np.allclose(q0[:128,:128], q1[:128,:128]))  #True
print(np.allclose(q0, q1))                        #False
```


Discussion:
This is because of an integer overflow in an internal indexing calculation. I'll send a fix shortly.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3879,6 +3879,10 @@ def replace(tup, val):
     lst[axis] = val
     return tuple(lst)
 
+  use_64bit_index = _any([type(d) is Poly or d >= (1 << 31) for d in arr.shape])
+  index_dtype = int64 if use_64bit_index else int32
+  indices = lax.convert_element_type(indices, index_dtype)
+
   bcast_shape = lax.broadcast_shapes(replace(arr.shape, 1), replace(indices.shape, 1))
   indices = broadcast_to(indices, replace(bcast_shape, indices.shape[axis]))
   arr     = broadcast_to(arr,     replace(bcast_shape, arr.shape[axis]))


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -3346,10 +3346,11 @@ def testTakeEmpty(self):
 
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_{}_ishape={}_axis={}".format(
-          jtu.format_shape_dtype_string(x_shape, dtype), i_shape, axis),
-       "rng_factory": rng_factory, "x_shape": x_shape, "i_shape": i_shape, "dtype": dtype,
-       "axis": axis}
+      {"testcase_name": "_{}_index={}_axis={}".format(
+          jtu.format_shape_dtype_string(x_shape, dtype),
+          jtu.format_shape_dtype_string(i_shape, index_dtype), axis),
+       "x_shape": x_shape, "i_shape": i_shape, "dtype": dtype,
+       "index_dtype": index_dtype, "axis": axis}
       for x_shape, i_shape in filter(
         _shapes_are_equal_length,
         filter(_shapes_are_broadcast_compatible,
@@ -3357,9 +3358,10 @@ def testTakeEmpty(self):
       for axis in itertools.chain(range(len(x_shape)), [-1],
                                   [cast(Optional[int], None)])
       for dtype in default_dtypes
-      for rng_factory in [jtu.rand_default]))
-  def testTakeAlongAxis(self, x_shape, i_shape, dtype, axis, rng_factory):
-    rng = rng_factory(self.rng())
+      for index_dtype in int_dtypes))
+  def testTakeAlongAxis(self, x_shape, i_shape, dtype, index_dtype, axis):
+    rng = jtu.rand_default(self.rng())
+
     i_shape = np.array(i_shape)
     if axis is None:
       i_shape = [np.prod(i_shape, dtype=np.int64)]
@@ -3370,7 +3372,11 @@ def testTakeAlongAxis(self, x_shape, i_shape, dtype, axis, rng_factory):
     def args_maker():
       x = rng(x_shape, dtype)
       n = np.prod(x_shape, dtype=np.int32) if axis is None else x_shape[axis]
-      i = rng(i_shape, np.int32) % (2 * n - 1) - (n - 1)
+      if np.issubdtype(index_dtype, np.unsignedinteger):
+        index_rng = jtu.rand_int(self.rng(), 0, n)
+      else:
+        index_rng = jtu.rand_int(self.rng(), -n, n)
+      i = index_rng(i_shape, index_dtype)
       return x, i
 
     jnp_op = lambda x, i: jnp.take_along_axis(x, i, axis=axis)
@@ -3380,6 +3386,14 @@ def args_maker():
       self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
     self._CompileAndCheck(jnp_op, args_maker)
 
+  def testTakeAlongAxisWithUint8IndicesDoesNotOverflow(self):
+    # https://github.com/google/jax/issues/5088
+    h = jtu.rand_default(self.rng())((256, 256, 100), np.float32)
+    g = jtu.rand_int(self.rng(), 0, 100)((256, 256, 1), np.uint8)
+    q0 = jnp.take_along_axis(h, g, axis=-1)
+    q1 = np.take_along_axis( h, g, axis=-1)
+    np.testing.assert_equal(q0, q1)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_n={}_increasing={}".format(
           jtu.format_shape_dtype_string([shape], dtype),