Repository: google/jax
Issue: Mysterious integer overflow in jax.vmap(jax.take)

Consider the following:
```python
>>> import jax.numpy as jnp
>>> indices = jnp.zeros(256, dtype=jnp.int8)
>>> choices = jnp.zeros(1, dtype=jnp.float32)
>>> choices[indices]
Array([0, 0, 0, 0, ...], dtype=float32)
>>> jnp.take(choices, indices, axis=0)
Array([0, 0, 0, 0, ...], dtype=float32)
```

Working as intended. But if we instead vmap the call to `jnp.take`:
```python
>>> jax.vmap(lambda x, i: jnp.take(x, i, axis=0))(jnp.stack([choices] * 256), indices)
Array([ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
       nan, nan, nan, nan, nan, nan, nan, nan, nan], dtype=float32)
```

Not working as intended! The `nan`s are being filled in as if there are out-of-bounds indices, even though all the values in the `int8` array are zero. Perhaps this code is causing some internal indexing calculations to be done in `int8`?

If you make the arrays longer, you can see the values oscillate between `0` and `nan`, which makes it clear it's some sort of overflow. See [this reproduction in Colab](https://colab.research.google.com/drive/1dRDM8z15x-Cu26i6TcLLFRPOu-8KT6H0?usp=sharing).

For what it's worth, the same issue seems to occur in `int16`, but only when you reach `2**31` in length. Overall strongly suggesting to me that something is being keyed off the index dtype, when it really shouldn't.

Also, it only seems to occur with `take(..., mode=None)` (the default) or `take(..., mode='fill')`, which both internally use `lax.GatherScatterMode.FILL_OR_DROP`. So perhaps the error can be found in the logic for filling. I don't know how to inspect the code for that, since it's in the `'gather'` primitive.

A workaround is to always upcast indices to (u)int32/int64, which is perhaps what is already happening with the regular indexing operation `x[i]`. I would also separately lobby for supporting mode 'promise_in_bounds' in `jnp.take`, and making that the default.

### What jax/jaxlib version are you using?

jax v0.4.11, jaxlib v0.4.11

### Which accelerator(s)?

Replicated on CPU, Nvidia GPU

Changed files:
jax/_src/lax/slicing.py, tests/lax_test.py

diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
index 7624ac071..e77f84028 100644
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -1329,10 +1329,12 @@ def _gather_batching_rule(batched_args, batch_dims, *, dimension_numbers,
     # indices shape (7, 3, 4, 5). We transform that to 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_dtype = _promote_dtype_for_size(indices.dtype, indices.shape[0])
     count_shape = list(indices.shape)
     count_shape[-1] = 1
-    counts = lax.broadcasted_iota(indices.dtype, tuple(count_shape), 0)
-    indices = lax.concatenate([counts, indices], len(count_shape) - 1)
+    counts = lax.broadcasted_iota(index_dtype, tuple(count_shape), 0)
+    indices = lax.concatenate([counts, indices.astype(index_dtype)],
+                              len(count_shape) - 1)
 
     slice_sizes = (1,) + slice_sizes
     collapsed_slice_dims = (0,) + tuple(np.add(1, dimension_numbers.collapsed_slice_dims))
@@ -1348,6 +1350,21 @@ def _gather_batching_rule(batched_args, batch_dims, *, dimension_numbers,
                   indices_are_sorted=indices_are_sorted, mode=mode,
                   fill_value=fill_value), 0
 
+def _promote_dtype_for_size(dtype, size):
+  if not dtypes.issubdtype(dtype, np.integer):
+    return dtype
+  # size may be a dynamic shape, in which case we return at least int32
+  try:
+    size = int(size)
+  except:
+    return dtype if np.iinfo(dtype).bits >= 32 else np.dtype('int32')
+  if size <= np.iinfo(dtype).max:
+    return dtype
+  elif size <= np.iinfo(np.int32).max:
+    return np.dtype('int32')
+  else:
+    return dtypes.canonicalize_dtype(np.int64)
+
 def _gather_pad_rule(in_avals, out_avals, operand, indices, *,
                      dimension_numbers, slice_sizes, unique_indices,
                      indices_are_sorted, mode, fill_value):
diff --git a/tests/lax_test.py b/tests/lax_test.py
index e5c753722..97d9ce4cf 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -3170,6 +3170,15 @@ class CustomElementTypesTest(jtu.JaxTestCase):
     self.assertIsInstance(ys, FooArray)
     self.assertEqual(ys.shape, (3, 2, 1))
 
+  def test_gather_batched_index_dtype(self):
+    # Regression test for https://github.com/google/jax/issues/16557
+    dtype = jnp.int8
+    size = jnp.iinfo(dtype).max + 10
+    indices = jnp.zeros(size, dtype=dtype)
+    values = jnp.zeros((size, 1))
+    results = jax.vmap(lambda x, i: jnp.take(x, i, axis=0))(values, indices)
+    self.assertArraysEqual(results, jnp.zeros(size))
+
   @parameterized.parameters([
     (0,),
     (slice(1),),