Repository: google/jax

Issue:
jax.numpy.take returns large negative number when given out of range index
### Description

If I use `jax.numpy.take` to try to retrieve an index out of bounds, I get a large negative number:
```python
import jax.numpy as jnp
arr = jnp.arange(5)
print(jnp.take(arr, 10)) # -2147483648
```

I know that out-of-bounds indexing doesn't error for performance/implementation reasons,  but is the value `-2147483648` really a great choice? I think it might be a bit unexpected, e.g. the [documentation](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#out-of-bounds-indexing) says:

> the index is clamped to the bounds of the array since something must be returned

But it seems that the index here isn't being clamped.

### What jax/jaxlib version are you using?

jax v0.4.2, jaxlib v0.4.2

### Which accelerator(s) are you using?

TPU

### Additional system info

Python 3.8.10, TPU VM v3-8

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report – this is expected behavior, though I understand it's somewhat unintuitive. There's some discussion of why JAX doesn' raise errors on out-of-bound indices here: http://go/jax-sharp-bits#out-of-bounds-indexing

So in this case, JAX must return some value of the specified dtype to indicate that the index was out-of-bounds. For float arrays, JAX will return `NaN`. But for int arrays, there is no equivalent of a `NaN` value, so for lack of a better choice JAX returns the minimum representable value. You can control this behavior using the `fill_value` and `mode` arguments, as discussed in the  [`jnp.take` docstring](http://go/jax-docs/_autosummary/jax.numpy.take.html):
> - mode (string, default="fill") – Out-of-bounds indexing mode. The default mode=”fill” returns invalid values (e.g. NaN) for out-of bounds indices. See [jax.numpy.ndarray.at](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at) for more discussion of out-of-bounds indexing in JAX.
> ...
> - fill_value (optional) – The fill value to return for out-of-bounds slices when mode is ‘fill’. Ignored otherwise. Defaults to NaN for inexact types, the largest negative value for signed types, the largest positive value for unsigned types, and True for booleans.
I see, thanks for the explanation. I missed those docs, as I just read the main docstring anticipating it would discuss the out-of-bounds behaviour. Maybe we could add something commenting on this behaviour in the main docstring? Something like the note in the docstring for [`jax.numpy.ndarray.at`](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at):

> By default, JAX assumes that all indices are in-bounds. There is experimental support for giving more precise semantics to out-of-bounds indexed accesses, via the mode parameter (see below).

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
@@ -3632,23 +3632,24 @@ def unpackbits(a, axis: Optional[int] = None, count=None, bitorder='big'):
 
 @_wraps(np.take, skip_params=['out'],
         lax_description="""
-The JAX version adds several extra parameters, described below, which are forwarded
-to :func:`jax.lax.gather` for finer control over indexing.""",
+By default, JAX assumes that all indices are in-bounds. Alternative out-of-bound
+index semantics can be specified via the ``mode`` parameter (see below).
+""",
         extra_params="""
 mode : string, default="fill"
     Out-of-bounds indexing mode. The default mode="fill" returns invalid values
-    (e.g. NaN) for out-of bounds indices. See :attr:`jax.numpy.ndarray.at` for
-    more discussion of out-of-bounds indexing in JAX.
+    (e.g. NaN) for out-of bounds indices (see also ``fill_value`` below).
+    For more discussion of mode options, see :attr:`jax.numpy.ndarray.at`.
+fill_value : optional
+    The fill value to return for out-of-bounds slices when mode is 'fill'. Ignored
+    otherwise. Defaults to NaN for inexact types, the largest negative value for
+    signed types, the largest positive value for unsigned types, and True for booleans.
 unique_indices : bool, default=False
     If True, the implementation will assume that the indices are unique,
     which can result in more efficient execution on some backends.
 indices_are_sorted : bool, default=False
     If True, the implementation will assume that the indices are sorted in
     ascending order, which can lead to more efficient execution on some backends.
-fill_value : optional
-    The fill value to return for out-of-bounds slices when mode is 'fill'. Ignored
-    otherwise. Defaults to NaN for inexact types, the largest negative value for
-    signed types, the largest positive value for unsigned types, and True for booleans.
 """)
 def take(a, indices, axis: Optional[int] = None, out=None, mode=None,
          unique_indices=False, indices_are_sorted=False, fill_value=None):
@@ -5168,9 +5169,8 @@ class _IndexUpdateHelper:
   in which conflicting updates are applied is implementation-defined and may be
   nondeterministic (e.g., due to concurrency on some hardware platforms).
 
-  By default, JAX assumes that all indices are in-bounds. There is experimental
-  support for giving more precise semantics to out-of-bounds indexed accesses,
-  via the ``mode`` parameter (see below).
+  By default, JAX assumes that all indices are in-bounds. Alternative out-of-bound
+  index semantics can be specified via the ``mode`` parameter (see below).
 
   Arguments
   ---------


Test patch: