internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: google/jax

Issue:
jnp.unique does not properly handle NaN values
Example:
```python
import numpy as np
import jax.numpy as jnp

x = np.array([0, 1, np.nan, np.nan])

print(np.unique(x))
# [ 0.,  1., nan])

print(jnp.unique(x))
[ 0.,  1., nan, nan]
```



Discussion:

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
@@ -5596,6 +5596,11 @@ def replace(tup, val):
 @partial(jit, static_argnums=1)
 def _unique_sorted_mask(ar, axis):
   aux = moveaxis(ar, axis, 0)
+  if issubdtype(aux.dtype, np.complexfloating):
+    # Work around issue in sorting of complex numbers with Nan only in the
+    # imaginary component. This can be removed if sorting in this situation
+    # is fixed to match numpy.
+    aux = where(isnan(aux), lax._const(aux, nan), aux)
   size, *out_shape = aux.shape
   if _prod(out_shape) == 0:
     size = 1
@@ -5604,7 +5609,13 @@ def _unique_sorted_mask(ar, axis):
     perm = lexsort(aux.reshape(size, _prod(out_shape)).T[::-1])
   aux = aux[perm]
   if aux.size:
-    mask = ones(size, dtype=bool).at[1:].set(any(aux[1:] != aux[:-1], tuple(range(1, aux.ndim))))
+    if issubdtype(aux.dtype, inexact):
+      # This is appropriate for both float and complex due to the documented behavior of np.unique:
+      # See https://github.com/numpy/numpy/blob/v1.22.0/numpy/lib/arraysetops.py#L212-L220
+      neq = lambda x, y: lax.ne(x, y) & ~(isnan(x) & isnan(y))
+    else:
+      neq = lax.ne
+    mask = ones(size, dtype=bool).at[1:].set(any(neq(aux[1:], aux[:-1]), tuple(range(1, aux.ndim))))
   else:
     mask = zeros(size, dtype=bool)
   return aux, mask, perm


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
@@ -2479,6 +2479,29 @@ def np_fun(x, fill_value=fill_value):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @unittest.skipIf(numpy_version < (1, 21), "Numpy < 1.21 does not properly handle NaN values in unique.")
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": f"_{dtype.__name__}", "dtype": dtype}
+      for dtype in inexact_dtypes))
+  def testUniqueNans(self, dtype):
+    def args_maker():
+      x = [-0.0, 0.0, 1.0, 1.0, np.nan, -np.nan]
+      if np.issubdtype(dtype, np.complexfloating):
+        x = [complex(i, j) for i, j in itertools.product(x, repeat=2)]
+      return [np.array(x, dtype=dtype)]
+
+    kwds = dict(return_index=True, return_inverse=True, return_counts=True)
+    jnp_fun = partial(jnp.unique, **kwds)
+    def np_fun(x):
+      dtype = x.dtype
+      # numpy unique fails for bfloat16 NaNs, so we cast to float64
+      if x.dtype == jnp.bfloat16:
+        x = x.astype('float64')
+      u, *rest = np.unique(x, **kwds)
+      return (u.astype(dtype), *rest)
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
+
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_fixed_size={}".format(fixed_size),
       "fixed_size": fixed_size}
Repository: google/jax

Issue:
[sparse] bug in bcoo_extract batching rule
```python
from jax.experimental import sparse
from jax import vmap
import jax.numpy as jnp

indices = jnp.zeros((4, 1, 1), dtype=int)
v = jnp.ones(1)

# Manual mapping
print(jnp.concatenate([sparse.bcoo_extract(ind, v) for ind in indices]))
# [1. 1. 1. 1.]

# vmap
vmap(sparse.bcoo_extract, in_axes=(0, None))(indices, v)
# ValueError: indices batch dimensions not compatible for indices.shape=(4, 1, 1), shape=(1, 1)
```


Discussion:

Resolving patch:
diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -426,10 +426,14 @@ def _bcoo_extract_batching_rule(batched_args, batch_dims):
     bdim = batch_dims[1]
     indices = lax.expand_dims(indices, (bdim,))
   elif batch_dims[1] is None:
+    # TODO(jakevdp) can we handle this case without explicit broadcasting?
     bdim = batch_dims[0]
-    mat = lax.expand_dims(mat, (bdim,))
+    result_shape = list(mat.shape)
+    result_shape.insert(bdim, indices.shape[bdim])
+    mat = lax.broadcast_in_dim(mat, result_shape, (bdim,))
   else:
-    assert batch_dims[0] == batch_dims[1]
+    if batch_dims[0] != batch_dims[1]:
+      raise NotImplementedError("bcoo_extract with unequal batch dimensions.")
     bdim = batch_dims[0]
   n_batch = indices.ndim - 2
   if bdim >= n_batch:


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -693,6 +693,26 @@ def test_bcoo_extract(self, shape, dtype, n_batch, n_dense):
     data3 = jit(sparse.bcoo_extract)(indices, M)
     self.assertArraysEqual(data, data3)
 
+  def test_bcoo_extract_batching(self):
+    # https://github.com/google/jax/issues/9431
+    indices = jnp.zeros((4, 1, 1), dtype=int)
+    mat = jnp.arange(4.).reshape((4, 1))
+
+    # in_axes = (0, None)
+    expected = jnp.vstack([sparse.bcoo_extract(i, mat[0]) for i in indices])
+    actual = vmap(sparse.bcoo_extract, in_axes=(0, None))(indices, mat[0])
+    self.assertArraysEqual(expected, actual)
+
+    # in_axes = (None, 0)
+    expected = jnp.vstack([sparse.bcoo_extract(indices[0], m) for m in mat])
+    actual = vmap(sparse.bcoo_extract, in_axes=(None, 0))(indices[0], mat)
+    self.assertArraysEqual(expected, actual)
+
+    # in_axes = (0, 0)
+    expected = jnp.vstack([sparse.bcoo_extract(i, m) for i, m in zip(indices, mat)])
+    actual = vmap(sparse.bcoo_extract, in_axes=0)(indices, mat)
+    self.assertArraysEqual(expected, actual)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_nbatch={}_ndense={}".format(
         jtu.format_shape_dtype_string(shape, dtype), n_batch, n_dense),
Repository: google/jax

Issue:
`DeviceArray.__mul__` etc. don't respect `__jax_array__`
```python
import jax.numpy as jnp

class M:
    def __jax_array__(self):
        return jnp.array([1, 1])

M() * jnp.array([1, 1])
```
produces `TypeError: unsupported operand type(s) for *: 'M' and 'DeviceArray'`

In contrast `jnp.multiply(M(), jnp.array([1, 1]))` handles this just fine. No error, works as expected.

Also, NumPy handles this fine:
```python
import numpy as np

class M:
    def __array__(self):
        return np.array([1, 1])

M() * np.array([1, 1])  # no error!
```


Discussion:
Thanks for the report. I think the issue is we need to check for `__jax_array__` here: https://github.com/google/jax/blob/f02c6fcd728878f8a239de895eb67836fb2f9b46/jax/_src/numpy/lax_numpy.py#L6777

I can work on a fix later this afternoon if nobody else gets to it first.

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
@@ -6774,6 +6774,8 @@ def wrapped(*args, **kwargs):
 def _defer_to_unrecognized_arg(binary_op):
   # Ensure that other array types have the chance to override arithmetic.
   def deferring_binary_op(self, other):
+    if hasattr(other, '__jax_array__'):
+      other = other.__jax_array__()
     if not isinstance(other, _accepted_binop_types):
       return NotImplemented
     return binary_op(self, other)


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
@@ -3852,6 +3852,12 @@ def __array__(self, dtype=None):
     ans = jnp.array(a)
     self.assertEqual(ans, 3.)
 
+  def testJaxArrayOps(self):
+    class arraylike:
+      def __jax_array__(self):
+        return jnp.array(3.)
+    self.assertArraysEqual(arraylike() * jnp.arange(10), jnp.array(3.) * jnp.arange(10))
+
   def testMemoryView(self):
     self.assertAllClose(
         jnp.array(bytearray(b'\x2a')),
Repository: google/jax

Issue:
jax.random.choice has undefined behavior when sampling without replacement and using probabilities
Summary: if enough weights/probabilities are zero such that sampling w/o replacement leads to too few possibilities to sample, jax should raise an error.
Simple case: 
```python
import jax
import jax.numpy as jnp

prob = jnp.zeros(10)
prob = prob.at[9].set(1)
jax.random.choice(jax.random.PRNGKey(5), jnp.arange(10), shape=(5,), replace=False, p=prob)
# DeviceArray([9, 0, 1, 2, 3], dtype=int32)
```
Numpy raises a 
```
ValueError: Fewer non-zero entries in p than size
```
with equivalent code.

Fix: 
https://github.com/google/jax/blob/main/jax/_src/random.py#L493
Here, do a check that the number of nonzero entries is >= desired number of entries. 



Discussion:
Sounds indeed like Jax should mimic the numpy behaviour, I can take a look.
Thanks for the report – the reason for this behavior is that JAX cannot raise warnings or errors that depend on the values within arrays. This is due to the nature of JAX's tracing and compilation. It's the same reason that, for example, we cannot raise errors for out-of-bound indexing.

Our best option here is to probably add a warning to the docstring, informing users that if there are not enough nonzero probabilities, the output is ill-defined. Are you interested in opening a pull request for that?
Interesting, sure I'll open new PR with a warning in the docstring. OK if I just close the other one?
> Interesting, sure I'll open new PR with a warning in the docstring. OK if I just close the other one?

Sure, sounds good.
So I've found that it is (mostly) possible to raise value-dependent errors under jit, by using `jax.experimental.host_callback`. For example see
https://github.com/patrick-kidger/diffrax/blob/2b4e4d863c15abc7143919bac7825090bbfe50be/diffrax/misc/errors.py#L15
for an implementation, and the comment on
#9457
for a way to get it (mostly) working on the GPU as well.

I find that this is really useful functionality. Maybe with some tidying-up this could find its way into `jax.experimental` somewhere?
Thanks for pointing that out – there's also work on value-dependent checks here: https://github.com/google/jax/blob/main/jax/experimental/checkify.py
Ah nice -- that's where #8468 ended up! Thanks for the link. Once the details of that have been worked out (being able to stage `checkify.check` + merging `checkify.checkify` into `jax.jit` (?)) then that looks like it'll be a much more elegant replacement.

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -435,6 +435,11 @@ def choice(key: KeyArray,
            axis: int = 0) -> jnp.ndarray:
   """Generates a random sample from a given array.
 
+  .. warning::
+    If ``p`` has fewer non-zero elements than the requested number of samples,
+    as specified in ``shape``, and ``replace=False``, the output of this
+    function is ill-defined. Please make sure to use appropriate inputs.
+
   Args:
     key: a PRNG key used as the random key.
     a : array or int. If an ndarray, a random sample is generated from


Test patch:
Repository: google/jax

Issue:
[QoL] Add copy button in docs code snippets
Since I'm a bit lazy, I'd like to have a "copy to clipboard" button in jax docs to copy over code snippets instead of drag-select-copying it.  Like this:

![image](https://user-images.githubusercontent.com/56438982/155105866-317c1ca8-5c1e-4809-be2a-2cf3a3145d23.png)

Dupicate Checks:
Nothing relevant comes up when searching for "copy button", "docs copy button" or even "button" for that matter.


Discussion:
After a bit of digging, I added a "sphinx-copybutton" to the requirements and `config.py`, built anew (without executing notebooks). And it seems to work. 

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -70,6 +70,7 @@
     'sphinx_autodoc_typehints',
     'myst_nb',
     "sphinx_remove_toctrees",
+    'sphinx_copybutton',
     'jax_extensions',
 ]
 


Test patch:
Repository: google/jax

Issue:
jax.random.poisson hangs if lam is NaN
The following causes Python to hang:

```python
import jax
import numpy as np

key = jax.random.PRNGKey(0)
jax.random.poisson(key=key, lam=jnp.array([np.nan]), shape=(1,))
# hangs
```

My expectation is that when lam=np.nan, the output will be np.nan.



Discussion:
Thanks for the report - it looks like this corner case leads to a very large (but not technically infinite) loop: https://github.com/google/jax/blob/main/jax/_src/random.py#L1104

We should be able to fix this pretty quickly.
The fix is in #9721. Note that we can't return NaN in this case, because the dtype of the returned array is integer. But the fix returns `-1`, similar to other invalid values of `lam`.
Wow, that was fast. Thanks!

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -53,6 +53,9 @@
 
 ### utilities
 
+def _isnan(x):
+  return lax.ne(x, x)
+
 def _check_prng_key(key):
   # TODO(frostig): remove once we always enable_custom_prng
   if type(key) is prng.PRNGKeyArray:
@@ -1096,7 +1099,7 @@ def _poisson(key, lam, shape, dtype):
   # https://github.com/numpy/numpy/blob/v1.18.3/numpy/random/src/distributions/distributions.c#L574
   # For lambda < 10, we use the Knuth algorithm; otherwise, we use transformed
   # rejection sampling.
-  use_knuth = lam < 10
+  use_knuth = _isnan(lam) | (lam < 10)
   lam_knuth = lax.select(use_knuth, lam, lax.full_like(lam, 0.0))
   # The acceptance probability for rejection sampling maxes out at 89% as
   # λ -> ∞, so pick some arbitrary large value.


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -791,6 +791,12 @@ def testPoissonZeros(self):
     samples = random.poisson(key, lam, shape=(2, 20))
     self.assertArraysEqual(samples[:, :10], jnp.zeros_like(samples[:, :10]))
 
+  def testPoissonCornerCases(self):
+    key = self.seed_prng(0)
+    lam = jnp.array([-1, 0, jnp.nan])
+    samples = random.poisson(key, lam, shape=(3,))
+    self.assertArraysEqual(samples, jnp.array([-1, 0, -1]))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}".format(np.dtype(dtype).name), "dtype": dtype}
       for dtype in jtu.dtypes.floating))
Repository: google/jax

Issue:
Add jax.scipy.linalg.eigh_tridiagonal to document
Just let users know this algorithm has been officially implemented.
I hope others can avoid wasting time like me:
> I check the document and do not find `eigh_tridiagonal`, thus I migrate tensorflow's implementation by myself. When I start to compare my code against `jax.scipy.linalg.eigh`, I suddenly find that there is `eigh_tridiagonal`!!(Thanks to vscode's auto-complete)😭😭😭😭


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -1461,7 +1461,7 @@ def _tridiagonal_solve_jax(dl, d, du, b, **kw):
 def tridiagonal_solve(dl, d, du, b):
   r"""Computes the solution of a tridiagonal linear system.
 
-  This function computes the solution of a tridiagonal linear system::
+  This function computes the solution of a tridiagonal linear system:
 
   .. math::
     A . X = B


Test patch:
Repository: google/jax

Issue:
ensure_compile_time_eval doc example fails
The example of [ensure_compile_time_eval](https://jax.readthedocs.io/en/latest/_autosummary/jax.ensure_compile_time_eval.html)

```python
import jax
import jax.numpy as jnp

print(jax.__version__)  # 0.3.1

@jax.jit
def f(x):
    with jax.ensure_compile_time_eval():
        y = jnp.sin(3.0)
        z = jnp.sin(y)
        #z = z.item()  # adding this works
    if z > 0:  # the value of z is availble and can be used in control flow
        return jnp.sin(x)
    else:
        return jnp.cos(x)
    
f(0)
```
gives ConcretizationTypeError.

Converting z to python float works.


Discussion:
Hah, oops, thanks! That makes sense. I'm embarrassed that I apparently didn't test this...

Also putting the `z > 0` inside the context manager works, i.e. having a line like `positive = z > 0` and then testing `if positive`. 

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -868,7 +868,8 @@ def f(x):
       with jax.ensure_compile_time_eval():
         y = jnp.sin(3.0)
         z = jnp.sin(y)
-      if z > 0:  # the value of z is availble and can be used in control flow
+        z_positive = z > 0
+      if z_positive:  # z_positive is usable in Python control flow
         return jnp.sin(x)
       else:
         return jnp.cos(x)


Test patch:
Repository: google/jax

Issue:
DOC: `disable_jit` scope of influence
Maybe the document of `disable_jit` should be updated.
> For debugging it is useful to have a mechanism that disables [jit()](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit) everywhere in a dynamic context.

The document says that `disable_jit` will disable `jit()`, not mentions that it will also affect behavior of control flow operators such as `lax.scan`.

_Originally posted by @YouJiacheng in https://github.com/google/jax/discussions/9913#discussioncomment-2373983_


Discussion:

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -544,7 +544,12 @@ def disable_jit():
   """Context manager that disables :py:func:`jit` behavior under its dynamic context.
 
   For debugging it is useful to have a mechanism that disables :py:func:`jit`
-  everywhere in a dynamic context.
+  everywhere in a dynamic context. Note that this not only disables explicit uses
+  of `jit` by the user, but will also remove any implicit JIT compilation used by the
+  JAX library: this includes implicit JIT computation of `body` and `cond`
+  functions passed to higher-level primitives like :func:`scan` and :func:`while_loop`,
+  JIT used in implementations of :mod:`jax.numpy` functions, and any other case where
+  `jit` is used within an API's implementation.
 
   Values that have a data dependence on the arguments to a jitted function are
   traced and abstracted. For example, an abstract value may be a


Test patch:
Repository: google/jax

Issue:
bug in jax.nn.glu
There is an issue with `jax.nn.glu`  in the latest release

https://github.com/google/jax/blob/728e4fd3fad334e551dcd1a13b47b8d51aae5a9f/jax/_src/nn/functions.py#L261

Here the `static_argnames` should be `axis`. Right now  calling `glu`  results in a `TracerIntegerConversionError`

```python
import jax
import jax.numpy as jnp

jax.nn.glu(jnp.ones((1, 4)), axis=1)
```

full error messages/tracebacks.
```
TracerIntegerConversionError              Traceback (most recent call last)
Input In [8], in <module>
----> 1 jax.nn.glu(jnp.ones((1, 4)), 1)

    [... skipping hidden 14 frame]

File ~/transformers-env/lib/python3.8/site-packages/jax/_src/nn/functions.py:269, in glu(x, axis)
    261 @partial(jax.jit, static_argnames=("glu",))
    262 def glu(x: Array, axis: int = -1) -> Array:
    263   """Gated linear unit activation function.
    264 
    265   Args:
    266     x : input array
    267     axis: the axis along which the split should be computed (default: -1)
    268   """
--> 269   size = x.shape[axis]
    270   assert size % 2 == 0, "axis size must be divisible by 2"
    271   x1, x2 = jnp.split(x, 2, axis)

File ~/transformers-env/lib/python3.8/site-packages/jax/core.py:473, in Tracer.__index__(self)
    472 def __index__(self):
--> 473   raise TracerIntegerConversionError(self)

TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError
```



Discussion:
Thanks for the report - it looks like a typo. The fact that this made it into a release makes me worried that we don't have good test coverage for the function.

I'll make a fix

Resolving patch:
diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -258,7 +258,7 @@ def gelu(x: Array, approximate: bool = True) -> Array:
   else:
     return jnp.array(x * (lax.erf(x / np.sqrt(2)) + 1) / 2, dtype=x.dtype)
 
-@partial(jax.jit, static_argnames=("glu",))
+@partial(jax.jit, static_argnames=("axis",))
 def glu(x: Array, axis: int = -1) -> Array:
   """Gated linear unit activation function.
 


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -81,7 +81,7 @@ def testEluValue(self):
     self.assertAllClose(val, 1e4, check_dtypes=False)
 
   def testGluValue(self):
-    val = nn.glu(jnp.array([1.0, 0.0]))
+    val = nn.glu(jnp.array([1.0, 0.0]), axis=0)
     self.assertAllClose(val, jnp.array([0.5]))
 
   @parameterized.parameters(False, True)
Repository: google/jax

Issue:
`jnp.take` out of bound documentation incorrect
Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

```
jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0
)
# DeviceArray([0, 1, 2, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4], dtype=int32)

jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0,
    mode="raise",
)
# NotImplementedError
```


According to the [documentation](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.take.html), the default out-of-bound behavior is "raise". However, this is not the case in practice


Discussion:
Thanks for the report - I'll work on a documentation fix.

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
@@ -3344,7 +3344,11 @@ def unpackbits(a, axis: Optional[int] = None, count=None, bitorder='big'):
   return swapaxes(unpacked, axis, -1)
 
 
-@_wraps(np.take, skip_params=['out'])
+@_wraps(np.take, skip_params=['out'], lax_description="""\
+In the JAX version, the ``mode`` argument defaults to a special version of ``"clip"``
+that handles negative indices. See :attr:`jax.numpy.ndarray.at` for more discussion
+of out-of-bounds indexing in JAX.
+""")
 def take(a, indices, axis: Optional[int] = None, out=None, mode=None):
   return _take(a, indices, None if axis is None else operator.index(axis), out,
                mode)


Test patch:
Repository: google/jax

Issue:
documentation format not correct
https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#:~:text=int32%2C%20weak_type%3DTrue)-,Summary,-%23

summary format not correct.
![3598734593485](https://user-images.githubusercontent.com/37238933/167804363-8ee6fec3-0bcf-4bc4-be0c-76bb7750a33e.png)




Discussion:
Thanks for the report – looks like some issue with myst-nb conversion. You can see the properly formatted version in the original notebook: https://colab.sandbox.google.com/github/google/jax/blob/main/docs/notebooks/Common_Gotchas_in_JAX.ipynb#scrollTo=SipXS5qiqk8e

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -184,6 +184,7 @@
 
 # -- Options for myst ----------------------------------------------
 myst_heading_anchors = 3  # auto-generate 3 levels of heading anchors
+myst_enable_extensions = ['dollarmath']
 nb_execution_mode = "force"
 nb_execution_allow_errors = False
 


Test patch: