anon-iclr-submission/benchname-bug-localization huggingface dataset:
Repository: google/jax
Issue: `repr` error on empty BCOO sparse matrix

I tried to initialize an empty BCOO matrix by:
```
from jax.experimental import sparse as jsparse
import jax.numpy as jnp

M = jsparse.BCOO((jnp.array([]), jnp.array([])), shape=(100, 100))
```

Just using the repr on `M` gives the following error: `<repr(<jax.experimental.sparse.ops.BCOO at 0x7f15fb4c9590>) failed: IndexError: tuple index out of range>`, which is unexpected.

On the other hand, initializing `M` with

```M = jsparse.BCOO((jnp.zeros(0, dtype='float32'), jnp.zeros((0, 2), dtype='int32')), shape=(100, 100)) ```

gives the correct repr output
```
M
BCOO(float32[100, 100], nse=0)
```



Changed files:
tests/sparse_test.py, jax/experimental/sparse/ops.py

diff --git a/jax/experimental/sparse/ops.py b/jax/experimental/sparse/ops.py
index 1b2c2c0e2..d6983ab04 100644
--- a/jax/experimental/sparse/ops.py
+++ b/jax/experimental/sparse/ops.py
@@ -1394,7 +1394,15 @@ class JAXSparse:
     self.shape = shape
 
   def __repr__(self):
-    repr_ = f"{self.__class__.__name__}({self.dtype}{list(self.shape)}, nse={self.nse})"
+    name = self.__class__.__name__
+    try:
+      nse = self.nse
+      dtype = self.dtype
+      shape = list(self.shape)
+    except:
+      repr_ = f"{name}(<invalid>)"
+    else:
+      repr_ = f"{name}({dtype}{shape}, nse={nse})"
     if isinstance(self.data, core.Tracer):
       repr_ = f"{type(self.data).__name__}[{repr_}]"
     return repr_
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
index e7592830b..060b44488 100644
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -1113,6 +1113,13 @@ class SparseGradTest(jtu.JaxTestCase):
 
 
 class SparseObjectTest(jtu.JaxTestCase):
+  def test_repr(self):
+    M = sparse.BCOO.fromdense(jnp.arange(5, dtype='float32'))
+    assert repr(M) == "BCOO(float32[5], nse=4)"
+
+    M_invalid = sparse.BCOO(([], []), shape=100)
+    assert repr(M_invalid) == "BCOO(<invalid>)"
+
   @parameterized.named_parameters(
     {"testcase_name": "_{}".format(Obj.__name__), "Obj": Obj}
     for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])
Repository: google/jax
Issue: [sparse] BCOO._dedupe moves padded indices to the front

Example:
```python
import jax.numpy as jnp
from jax.experimental import sparse

data = jnp.array([1, 1])
indices = jnp.array([1, 1])[:, None]
shape = (2,)

x = sparse.BCOO((data, indices), shape=shape)
x = x._dedupe()
print(x.data, x.indices.ravel())
# [2 0] [1 0]

x = x._dedupe()
print(x.data, x.indices.ravel())
# [0 2] [0 1]
```
The second `_dedupe()` call should leave the representation unchanged.

Changed files:
tests/sparse_test.py, jax/experimental/sparse/bcoo.py

diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
index e21a09508..11bd06875 100644
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -67,17 +67,20 @@ def _dedupe_bcoo_one(data, indices):
   if indices.shape[1] == 0:
     return data, indices
 
-  # This is a fixed-size version of jnp.unique() with return_indices=True
-  # unique values are zero-filled at the end.
-  perm = jnp.lexsort(indices.T[::-1])
+  # The following is similar to
+  #   indices_unique, inv_idx = jnp.unique(indices, axis=0, return_inverse=True,
+  #                                        size=indices.shape[0], fill_value=0)
+  # but modified to keep padding at the end of the resulting arrays.
+  is_padding = (indices == 0).all(1) & (data == 0)
+  perm = jnp.lexsort(indices[:, ::-1].T)
   aux = indices[perm]
   mask = jnp.ones(indices.shape[0], dtype=bool)
   mask = mask.at[1:].set(jnp.any(aux[1:] != aux[:-1], 1))
+  mask = mask & ~is_padding[perm]  # this is the padding modification.
   imask = jnp.cumsum(mask) - 1
   indices_unique = jnp.where(mask[:, None], aux, 0)[jnp.argsort(~mask)]
   inv_idx = jnp.zeros_like(imask).at[perm].set(imask)
 
-  # With the above, de-duping is easy.
   data_unique = jnp.zeros_like(data).at[inv_idx].add(data)
   return data_unique, indices_unique
 
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
index 87e93de96..8110daf03 100644
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -1059,6 +1059,16 @@ class BCOOTest(jtu.JaxTestCase):
     M_dedup = M._dedupe()
     self.assertAllClose(M.todense(), M_dedup.todense())
 
+  def test_bcoo_dedupe_padding(self):
+    # Regression test for https://github.com/google/jax/issues/8163
+    data = jnp.array([1, 0, 0])
+    indices = jnp.array([1, 0, 0])[:, None]
+    x = sparse.BCOO((data, indices), shape=(3,))
+    y = x._dedupe()
+    self.assertArraysEqual(x.todense(), y.todense())
+    self.assertArraysEqual(x.indices, y.indices)
+    self.assertArraysEqual(x.data, y.data)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_nbatch={}_ndense={}_axes={}".format(
         jtu.format_shape_dtype_string(shape, dtype), n_batch, n_dense, axes),
Repository: google/jax
Issue: vmap documentation does not document "axis_name" parameter

vmap() accepts an axis_name parameter (as verified by looking at the code), but this parameter is not documented at all in the docstring.

(From visiting https://jax.readthedocs.io/en/latest/_modules/jax/_src/api.html#vmap just prior to filing this bug)



Changed files:
jax/_src/api.py

diff --git a/jax/_src/api.py b/jax/_src/api.py
index 757dd7f9f..53b883bf2 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1329,6 +1329,8 @@ def vmap(fun: F, in_axes=0, out_axes=0, axis_name=None) -> F:
       (axes) of the array returned by the :func:`vmap`-ed function, which is one
       more than the number of dimensions (axes) of the corresponding array
       returned by ``fun``.
+    axis_name: Optional, a hashable Python object used to identify the mapped
+      axis so that parallel collectives can be applied.
 
   Returns:
     Batched/vectorized version of ``fun`` with arguments that correspond to
@@ -1403,6 +1405,16 @@ def vmap(fun: F, in_axes=0, out_axes=0, axis_name=None) -> F:
 
   If the ``out_axes`` is specified for a mapped result, the result is transposed
   accordingly.
+
+  Finally, here's an example using ``axis_name`` together with collectives:
+
+  >>> xs = jnp.arange(3. * 4.).reshape(3, 4)
+  >>> print(vmap(lambda x: lax.psum(x, 'i'), axis_name='i')(xs))
+  [[12. 15. 18. 21.]
+   [12. 15. 18. 21.]
+   [12. 15. 18. 21.]]
+
+  See the :py:func:`jax.pmap` docstring for more examples involving collectives.
   """
   _check_callable(fun)
   docstr = ("Vectorized version of {fun}. Takes similar arguments as {fun} "
Repository: google/jax
Issue: [sparse] sparsify transform does not respect weak types

```python
from jax import config
config.update('jax_enable_x64', True)

import jax.numpy as jnp
from jax.experimental import sparse

M = jnp.arange(12, dtype='int32').reshape(3, 4)
Msp = sparse.BCOO.fromdense(M)

print((2 * M).dtype)  # int32
print((2 * Msp).dtype)  # int64
```
Found in the process of #8180

Changed files:
tests/sparsify_test.py, jax/experimental/sparse/transform.py

diff --git a/jax/experimental/sparse/transform.py b/jax/experimental/sparse/transform.py
index 5b806a983..1f2ef7e07 100644
--- a/jax/experimental/sparse/transform.py
+++ b/jax/experimental/sparse/transform.py
@@ -55,6 +55,7 @@ import numpy as np
 from jax import core
 from jax import lax
 from jax import linear_util as lu
+import jax.numpy as jnp
 from jax._src.api_util import flatten_fun_nokwargs
 from jax.interpreters import partial_eval as pe
 from jax.interpreters import xla
@@ -78,7 +79,7 @@ class SparseEnv:
     self._buffers = list(bufs)
 
   def push(self, arr: Array) -> int:
-    self._buffers.append(np.array(arr) if np.isscalar(arr) else arr)  # type: ignore
+    self._buffers.append(jnp.asarray(arr))  # type: ignore
     return len(self._buffers) - 1
 
   def get(self, ind: int) -> Array:
@@ -155,7 +156,8 @@ def argspecs_to_avals(
     if argspec.is_unit():
       return core.abstract_unit
     else:
-      return core.ShapedArray(argspec.shape, argspec.data(spenv).dtype)
+      data = argspec.data(spenv)
+      return core.ShapedArray(argspec.shape, data.dtype, data.aval.weak_type)
   return tree_map(argspec_to_aval, argspecs, is_leaf=_is_argspec)
 
 
diff --git a/tests/sparsify_test.py b/tests/sparsify_test.py
index 2bcedf1e0..420305b16 100644
--- a/tests/sparsify_test.py
+++ b/tests/sparsify_test.py
@@ -343,6 +343,16 @@ class SparsifyTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(TypeError, "sparsified true_fun and false_fun output.*"):
       func(x_bcoo, y)
 
+  def testWeakTypes(self):
+    # Regression test for https://github.com/google/jax/issues/8267
+    M = jnp.arange(12, dtype='int32').reshape(3, 4)
+    Msp = BCOO.fromdense(M)
+    self.assertArraysEqual(
+      operator.mul(2, M),
+      sparsify(operator.mul)(2, Msp).todense(),
+      check_dtypes=True,
+    )
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax
Issue: two-operand cond conflicts with per-branch-operand cond

#8467 introduces multi-operand `lax.cond`. When given two operands, as in:

```python
def f(x):
  return lax.cond(x == 0, lax.add, lax.mul, x, y)
```

`cond` detects five total arguments and falls back to behavior that is backwards-compatible with the deprecated per-branch-operand `cond`. The latter accepts, in order: a predicate, true operand, true branch, false operand, and false branch. This results in an error:

```
TypeError: Value <function add ...> with type <class 'function'> is not a valid JAX type
```

One option is to rollback (or disallow) multi-operand cond for now, complete the deprecation of the previous cond, and roll forward (or re-enable) multi-operand cond.

Changed files:
tests/lax_control_flow_test.py, jax/_src/lax/control_flow.py

diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
index 2f8c92981..1d75f4a97 100644
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -788,7 +788,10 @@ def cond(*args, **kwargs):
   except TypeError:
     pass
   else:
-    return _cond_with_per_branch_args(*ba.args)
+    assert not ba.kwargs  # no catch-all **kwargs in _cond_with_per_branch
+    _, _, maybe_true_fun, _, maybe_false_fun = ba.args
+    if callable(maybe_true_fun) and callable(maybe_false_fun):
+      return _cond_with_per_branch_args(*ba.args)
 
   return _cond(*args, **kwargs)
 
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index 5a703d774..3a79710f8 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -599,8 +599,6 @@ class LaxControlFlowTest(jtu.JaxTestCase):
 
   def testCondTwoOperands(self):
     # see https://github.com/google/jax/issues/8469
-    self.skipTest("two-operand cond behavior is ambiguous (#8469)")
-
     add, mul = lax.add, lax.mul
 
     def fun(x):
Repository: google/jax
Issue: Unexpected assertion error for gradient of ppermute.

The following code produces an assertion error:

```
import jax

def ppermute(input):
  return jax.lax.ppermute(input, axis_name="i", perm=[[0, 1], [1, 0]])

grad_fn = jax.grad(ppermute)

vmapped_gradients_fn = jax.vmap(grad_fn, axis_name="i")
  
vector = jax.numpy.array([1., 2.])
grads = vmapped_gradients_fn(vector)
```

```
AssertionError                            Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/jax/_src/lax/parallel.py in _ppermute_batcher(axis_size, frame_name, _, vals_in, dims_in, axis_name, perm)
    791   assert axis_name[0] == frame_name, "ppermute batcher called with a wrong axis!"
    792   assert len(perm) == axis_size, "Permutation doesn't match the axis size!"
--> 793   assert d is not batching.not_mapped
    794   perm_indices = np.zeros(axis_size, dtype=int)
    795   for src, dst in perm:
```


Interestingly, the assertion error goes away, if I just inject a constant vector of ones:

```
import jax

def ppermute(input, one):
  return one * jax.lax.ppermute(input, axis_name="i", perm=[[0, 1], [1, 0]])

grad_fn = jax.grad(ppermute)

vmapped_gradients_fn = jax.vmap(grad_fn, axis_name="i")
  
vector = jax.numpy.array([1., 2.])
ones = jax.numpy.ones([2])
grads = vmapped_gradients_fn(vector, ones)
print(grads)  # [1. 1.]
```

This makes me believe the assert triggering above may indicate a bug.


Changed files:
jax/_src/lax/parallel.py, tests/batching_test.py

diff --git a/jax/_src/lax/parallel.py b/jax/_src/lax/parallel.py
index cd2e04c28..6a123f721 100644
--- a/jax/_src/lax/parallel.py
+++ b/jax/_src/lax/parallel.py
@@ -793,7 +793,8 @@ def _ppermute_batcher(axis_size, frame_name, _, vals_in, dims_in, axis_name, per
     raise NotImplementedError("ppermute batcher only supports a single axis")
   assert axis_name[0] == frame_name, "ppermute batcher called with a wrong axis!"
   assert len(perm) == axis_size, "Permutation doesn't match the axis size!"
-  assert d is not batching.not_mapped
+  if d is batching.not_mapped:
+    return v, d
   perm_indices = np.zeros(axis_size, dtype=int)
   for src, dst in perm:
     perm_indices[dst] = src
diff --git a/tests/batching_test.py b/tests/batching_test.py
index e822a9d06..bb6f2a72d 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -1267,6 +1267,19 @@ class BatchingTest(jtu.JaxTestCase):
     self.assertEqual(f(jnp.ones(3)).shape, (3,))
     self.assertEqual(jax.vmap(f)(jnp.ones((2, 3))).shape, (2, 3))
 
+  def testPpermuteBatcherTrivial(self):
+    # https://github.com/google/jax/issues/8688
+    def ppermute(input):
+      return jax.lax.ppermute(input, axis_name="i", perm=[[0, 1], [1, 0]])
+
+    grad_fn = jax.grad(ppermute)
+
+    vmapped_gradients_fn = jax.vmap(grad_fn, axis_name="i")
+
+    vector = jax.numpy.array([1., 2.])
+    ans = vmapped_gradients_fn(vector)  # doesn't crash
+    self.assertAllClose(ans, jnp.ones(2), check_dtypes=False)
+
 
 Array = Any
 ArrayElt = Any
Repository: google/jax
Issue: 0.2.9 fails: ModuleNotFoundError: No module named 'jaxlib'

```
Python 3.7.9 (default, Dec 12 2020, 01:25:17) 
[Clang 8.0.1 (tags/RELEASE_801/final 366581)] on freebsd12
Type "help", "copyright", "credits" or "license" for more information.
```

These lines:
```
import jax.numpy as jnp
from jax import grad, jit, vmap
```
fail:
```
$ python3.7 j.py 
Traceback (most recent call last):
  File "j.py", line 1, in <module>
    import jax.numpy as jnp
  File "/usr/local/lib/python3.7/site-packages/jax/__init__.py", line 22, in <module>
    from .api import (
  File "/usr/local/lib/python3.7/site-packages/jax/api.py", line 39, in <module>
    from . import core
  File "/usr/local/lib/python3.7/site-packages/jax/core.py", line 31, in <module>
    from . import dtypes
  File "/usr/local/lib/python3.7/site-packages/jax/dtypes.py", line 32, in <module>
    from .lib import xla_client
  File "/usr/local/lib/python3.7/site-packages/jax/lib/__init__.py", line 23, in <module>
    import jaxlib
ModuleNotFoundError: No module named 'jaxlib'
```

installed from the FreeBSD port
OS: FreeBSD 12.2

Downstream bug report: https://bugs.freebsd.org/bugzilla/show_bug.cgi?id=252871

Changed files:
jax/lib/__init__.py

diff --git a/jax/lib/__init__.py b/jax/lib/__init__.py
index e3a2cd8f1..894e6a785 100644
--- a/jax/lib/__init__.py
+++ b/jax/lib/__init__.py
@@ -20,7 +20,13 @@ __all__ = [
   'pocketfft', 'pytree', 'tpu_client', 'version', 'xla_client'
 ]
 
-import jaxlib
+try:
+  import jaxlib
+except ModuleNotFoundError as err:
+  raise ModuleNotFoundError(
+    'jax requires jaxlib to be installed. See '
+    'https://github.com/google/jax#installation for installation instructions.'
+    ) from err
 
 # Must be kept in sync with the jaxlib version in build/test-requirements.txt
 _minimum_jaxlib_version = (0, 1, 60)
Repository: google/jax
Issue: Unraveling an array with dtype bool results in pytree arrays with dtype float32

Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

```python
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
x = jnp.arange(10, dtype=jnp.float32)
x_flat, unravel = ravel_pytree(x)
y = x_flat < 5.3
print(y.dtype)          # => <dtype: 'bool'>
print(unravel(y).dtype) # => <dtype: 'float32'>
```

- [x] If applicable, include full error messages/tracebacks.
n/a

I'm running Python 3.9 with jax 0.2.19 and jaxlib 0.1.70.

Changed files:
tests/tree_util_test.py, jax/_src/flatten_util.py

diff --git a/jax/_src/flatten_util.py b/jax/_src/flatten_util.py
index 2f7bce645..b7d7dc5b8 100644
--- a/jax/_src/flatten_util.py
+++ b/jax/_src/flatten_util.py
@@ -57,7 +57,23 @@ def _ravel_list(lst):
   sizes, shapes = unzip2((jnp.size(x), jnp.shape(x)) for x in lst)
   indices = np.cumsum(sizes)
 
+  if all(dt == to_dtype for dt in from_dtypes):
+    # Skip any dtype conversion, resulting in a dtype-polymorphic `unravel`.
+    # See https://github.com/google/jax/issues/7809.
+    del from_dtypes, to_dtype
+    def unravel(arr):
+      chunks = jnp.split(arr, indices[:-1])
+      return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]
+    raveled = jnp.concatenate([jnp.ravel(e) for e in lst])
+    return raveled, unravel
+
+  # When there is more than one distinct input dtype, we perform type
+  # conversions and produce a dtype-specific unravel function.
   def unravel(arr):
+    arr_dtype = dtypes.dtype(arr)
+    if arr_dtype != to_dtype:
+      raise TypeError(f"unravel function given array of dtype {arr_dtype}, "
+                      f"but expected dtype {to_dtype}")
     chunks = jnp.split(arr, indices[:-1])
     with warnings.catch_warnings():
       warnings.simplefilter("ignore")  # ignore complex-to-real cast warning
diff --git a/tests/tree_util_test.py b/tests/tree_util_test.py
index 0164f6d8a..0fd2ded31 100644
--- a/tests/tree_util_test.py
+++ b/tests/tree_util_test.py
@@ -379,6 +379,23 @@ class RavelUtilTest(jtu.JaxTestCase):
     tree_ = unravel(raveled)
     self.assertAllClose(tree, tree_, atol=0., rtol=0.)
 
+  def testDtypePolymorphicUnravel(self):
+    # https://github.com/google/jax/issues/7809
+    x = jnp.arange(10, dtype=jnp.float32)
+    x_flat, unravel = flatten_util.ravel_pytree(x)
+    y = x_flat < 5.3
+    x_ = unravel(y)
+    self.assertEqual(x_.dtype, y.dtype)
+
+  def testDtypeMonomorphicUnravel(self):
+    # https://github.com/google/jax/issues/7809
+    x1 = jnp.arange(10, dtype=jnp.float32)
+    x2 = jnp.arange(10, dtype=jnp.int32)
+    x_flat, unravel = flatten_util.ravel_pytree((x1, x2))
+    y = x_flat < 5.3
+    with self.assertRaisesRegex(TypeError, 'but expected dtype'):
+      _ = unravel(y)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax
Issue: dynamic_update_slice start index adjustment fails when vmapped

The documentation for `dynamic_update_slice` explains "If the update slice is too large to fit in the array, the start index will be adjusted to make it fit" and this behavior does indeed seem to work normally but fails when I apply vmap. No update occurs if the update slice would not fit even if the update slice would still be partially in the array as in the example below.

```python
zeros = jnp.zeros(5)
ones = jnp.ones(3)

for i in range(1, 4):
    print(jax.lax.dynamic_update_slice(zeros, ones, (jnp.array(i, dtype=int),)))

print(jax.vmap(jax.lax.dynamic_update_slice, (None, None, 0))(zeros, ones, (jnp.arange(1, 4),)))

# [0. 1. 1. 1. 0.]
# [0. 0. 1. 1. 1.]
# [0. 0. 1. 1. 1.]
# [[0. 1. 1. 1. 0.]
#  [0. 0. 1. 1. 1.]
#  [0. 0. 0. 0. 0.]]
```

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

diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
index 9a5873739..5dc40ea4d 100644
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -933,7 +933,7 @@ def _dynamic_update_slice_batching_rule(batched_args, batch_dims):
     scatter, (operand, index, update), (operand_bd, index_bdim, update_bd),
     update_jaxpr=None, update_consts=None, dimension_numbers=dnums,
     indices_are_sorted=True, unique_indices=True,
-    mode=GatherScatterMode.PROMISE_IN_BOUNDS)
+    mode=GatherScatterMode.CLIP)
 
 
 dynamic_update_slice_p = standard_primitive(
diff --git a/jax/experimental/jax2tf/tests/shape_poly_test.py b/jax/experimental/jax2tf/tests/shape_poly_test.py
index 6c5d73380..2b8640c37 100644
--- a/jax/experimental/jax2tf/tests/shape_poly_test.py
+++ b/jax/experimental/jax2tf/tests/shape_poly_test.py
@@ -1701,6 +1701,9 @@ def _add_vmap_primitive_harnesses():
       "lu",
       "custom_linear_solve",
 
+      # Broken by https://github.com/google/jax/pull/9094
+      "dynamic_update_slice",
+
       # We do *= shapes in the batching rule for conv_general_dilated
       "conv_general_dilated",
 
diff --git a/tests/lax_test.py b/tests/lax_test.py
index f785fddc9..1a8ec5da2 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1552,6 +1552,15 @@ class LaxTest(jtu.JaxTestCase):
     self._CheckAgainstNumpy(lax_reference.dynamic_update_slice,
                             lax.dynamic_update_slice, args_maker)
 
+  def testDynamicUpdateSliceBatched(self):
+    # Regression test for https://github.com/google/jax/issues/9083
+    x = jnp.arange(5)
+    y = jnp.arange(6, 9)
+    ind = jnp.arange(6)
+    expected = jnp.vstack([lax.dynamic_update_slice(x, y, (i,)) for i in ind])
+    actual = jax.vmap(lax.dynamic_update_slice, (None, None, 0))(x, y, (ind,))
+    self.assertAllClose(expected, actual)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_perm={}".format(
           jtu.format_shape_dtype_string(shape, dtype), perm),
Repository: google/jax
Issue: Jax numpy array_split is slow compared to numpy

I was writing Jax code, substituting all of the places where I would normally use NumPy with jax.numpy
In using jax.numpy.array_split with a large input, I found that the program was incredibly slow (for jax+numpy or jax then the program didn't finish in 5 minutes for me) 

While running there was no issue however if I ctrl+c to terminate then the follow error was thrown
`F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:446] ptxas returned an error during compilation of ptx to sass: 'INTERNAL: ptxas exited with non-zero error code 2, output: '  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
Aborted (core dumped)`

It seems that Jax is trying to jit/compile the array input which is slow due to the size, I just wanted to check if this was expected or not
If this is expected, could some documentation be added to note this problem on the relative functions 

Example code
```
from tqdm import tqdm
import numpy as onp
import jax.numpy as jnp

print('numpy')
a = [batch for batch in tqdm(onp.array_split(onp.arange(10_000), 1_000))]
b = [batch for batch in tqdm(onp.array_split(onp.arange(100_000), 10_000))]

print('jax + numpy')
c = [batch for batch in tqdm(jnp.array_split(onp.arange(100_000), 10_000))]
d = [batch for batch in tqdm(jnp.array_split(onp.arange(100_000), 10_000))]

print('jax')
e = [batch for batch in tqdm(jnp.array_split(jnp.arange(100_000), 10_000))]
f = [batch for batch in tqdm(jnp.array_split(jnp.arange(100_000), 10_000))]
```

Thanks for any help


Changed files:
jax/_src/numpy/lax_numpy.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 37d27140e..6e4e44b80 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2260,6 +2260,8 @@ def broadcast_to(arr, shape):
 
 
 def _split(op, ary, indices_or_sections, axis=0):
+  _check_arraylike(op, ary)
+  ary = asarray(ary)
   axis = core.concrete_or_error(int, axis, f"in jax.numpy.{op} argument `axis`")
   size = ary.shape[axis]
   if isinstance(indices_or_sections, (tuple, list)):
@@ -2298,15 +2300,15 @@ def _split(op, ary, indices_or_sections, axis=0):
 def split(ary, indices_or_sections, axis: int = 0):
   return _split("split", ary, indices_or_sections, axis=axis)
 
-def _split_on_axis(np_fun, axis):
-  @_wraps(np_fun, update_doc=False)
+def _split_on_axis(op, axis):
+  @_wraps(getattr(np, op), update_doc=False)
   def f(ary, indices_or_sections):
-    return split(ary, indices_or_sections, axis=axis)
+    return _split(op, ary, indices_or_sections, axis=axis)
   return f
 
-vsplit = _split_on_axis(np.vsplit, axis=0)
-hsplit = _split_on_axis(np.hsplit, axis=1)
-dsplit = _split_on_axis(np.dsplit, axis=2)
+vsplit = _split_on_axis("vsplit", axis=0)
+hsplit = _split_on_axis("hsplit", axis=1)
+dsplit = _split_on_axis("dsplit", axis=2)
 
 @_wraps(np.array_split)
 def array_split(ary, indices_or_sections, axis: int = 0):
Repository: google/jax
Issue: Various type inconsistencies

Please consider adding `jax.numpy.generic` as an alias for `numpy.generic` (it seems to be missing).  Also,
```
In [5]: jnp.issubdtype(jnp.bfloat16, jnp.inexact)
Out[5]: True

In [6]: issubclass(jnp.bfloat16, jnp.inexact)
Out[6]: False

In [7]: issubclass(np.float32, np.inexact)
Out[7]: True
```
Line 6 should probably return true, which means somehow registering the types like numpy does.

Changed files:
jax/_src/numpy/lax_numpy.py, jax/numpy/__init__.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 6e4e44b80..979478513 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -424,6 +424,7 @@ uint = uint32 if dtypes.uint == np.uint32 else uint64
 float_ = float32 if dtypes.float_ == np.float32 else float64
 complex_ = complex64 if dtypes.complex_ == np.complex64 else complex128
 
+generic = np.generic
 number = np.number
 inexact = np.inexact
 complexfloating = np.complexfloating
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
index f28ec37c5..d51dc6a19 100644
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -160,6 +160,7 @@ from jax._src.numpy.lax_numpy import (
     full as full,
     full_like as full_like,
     gcd as gcd,
+    generic as generic,
     geomspace as geomspace,
     get_printoptions as get_printoptions,
     gradient as gradient,
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!
```

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 9e195f981..c504602af 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -6774,6 +6774,8 @@ _accepted_binop_types = (int, float, complex, np.generic, np.ndarray, ndarray)
 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)
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index e8225b1f6..1c0941789 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -3852,6 +3852,12 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     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: TypeError: No constant handler for type: <class 'jax.core.Token'>

Details can be found here.
https://github.com/mpi4jax/mpi4jax/pull/111


Changed files:
jax/interpreters/mlir.py, jax/interpreters/xla.py, tests/api_test.py

diff --git a/jax/interpreters/mlir.py b/jax/interpreters/mlir.py
index b7aafaf07..771ad7d4e 100644
--- a/jax/interpreters/mlir.py
+++ b/jax/interpreters/mlir.py
@@ -258,6 +258,8 @@ def _device_array_constant_handler(val, canonicalize_types):
 for t in device_array.device_array_types:
   register_constant_handler(t, _device_array_constant_handler)
 
+register_constant_handler(
+  core.Token, lambda _, __: [mhlo.CreateTokenOp(mhlo.TokenType.get())])
 
 # Source locations
 
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 7ab78d7df..606418958 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -374,6 +374,8 @@ for t in device_array.device_array_types:
   register_constant_handler(t, _device_array_constant_handler)
 
 
+register_constant_handler(core.Token, lambda c, _, __: [xops.CreateToken(c)])
+
 # TODO(mattjj): try to remove this canonicalize_dtype stuff
 def canonicalize_dtype(x):
   typ = type(x)
diff --git a/tests/api_test.py b/tests/api_test.py
index 78a33c3e9..c4d012b5a 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2959,6 +2959,11 @@ class APITest(jtu.JaxTestCase):
     x = jax.jit(jax.lax.create_token)(1.0)
     self.assertIsInstance(x, jax.core.Token)
 
+  def test_jit_capturing_token(self):
+    tok = jax.core.token
+    _, y = jax.jit(lambda x: (x + 2, tok))(7)
+    self.assertIsInstance(y, jax.core.Token)
+
   def test_leak_checker_catches_a_jit_leak(self):
     with jax.checking_leaks():
       lst = []
Repository: google/jax
Issue: Jax PRNGKeyArray.shape is a method, not a property

Was this an intended change? 

```python
>>> import jax
>>> jax.config.update('jax_enable_custom_prng', True)
>>> k=jax.random.PRNGKey(0)
>>> k.shape
<bound method PRNGKeyArray.shape of PRNGKeyArray:
  shape = ()
  impl = PRNGImpl:
           key_shape = (2,)
           seed = <function threefry_seed at 0x12e6f00d0>
           split = <function threefry_split at 0x10cc2aaf0>
           random_bits = <CompiledFunction object at 0x12e6c8880>
           fold_in = <function threefry_fold_in at 0x10cc2a790>>
>>> k.shape()
()
```

Changed files:
jax/_src/prng.py, tests/random_test.py

diff --git a/jax/_src/prng.py b/jax/_src/prng.py
index bfe498dd3..cde45b2b2 100644
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -125,6 +125,7 @@ class PRNGKeyArray:
           'deprecated `dtype` attribute of PRNG key arrays', FutureWarning)
       return np.uint32
 
+  @property
   def shape(self):
     # TODO(frostig): simplify once we always enable_custom_prng
     if config.jax_enable_custom_prng:
diff --git a/tests/random_test.py b/tests/random_test.py
index 3c0bf858e..7cc59d498 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1095,6 +1095,22 @@ class LaxRandomWithCustomPRNGTest(LaxRandomTest):
   def seed_prng(self, seed):
     return prng.seed_with_impl(double_threefry_prng_impl, seed)
 
+  def test_split_shape(self):
+    key = self.seed_prng(73)
+    keys = random.split(key, 10)
+    self.assertEqual(keys.shape, (10,))
+
+  def test_vmap_fold_in_shape(self):
+    key = self.seed_prng(73)
+    keys = vmap(lambda i: random.fold_in(key, i))(jnp.arange(3))
+    self.assertEqual(keys.shape, (3,))
+
+  def test_cannot_add(self):
+    key = self.seed_prng(73)
+    self.assertRaisesRegex(
+        TypeError, r'unsupported operand type\\(s\\) for \\+*',
+        lambda: key + 47)
+
 def _sampler_unimplemented_with_custom_prng(*args, **kwargs):
   raise SkipTest('sampler only implemented for default RNG')
Repository: google/jax
Issue: Keyword arguments to jnp.ndarray.clip() do not follow the documentation

The following should work according to the docs (https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.html), and it does work when using `numpy` instead of `jax.numpy`:

```python
import jax.numpy as jnp
jnp.zeros(3).clip(min=1)
```

However, in contrast to the docs, `jnp.ndarray.clip` actually expects the arguments to be called `a_min` and `a_max` respectively, just like `jnp.clip` does – i.e. `jnp.zeros(3).clip(a_min=1)` works.

Changed files:
jax/_src/numpy/lax_numpy.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 3dac38711..6a93b16e9 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -5946,6 +5946,20 @@ def _nbytes(arr):
   return size(arr) * _dtype(arr).itemsize
 
 
+def _clip(number, min=None, max=None, out=None, *, a_min=None, a_max=None):
+  # ndarray.clip has a slightly different API from clip (min -> a_min, max -> a_max)
+  # TODO: remove after deprecation window
+  if a_min is not None or a_max is not None:
+    warnings.warn('`a_min` and `a_max` keyword arguments to ndarray.clip are deprecated '
+                  'in favor of `min` and `max` for compatibility with numpy. '
+                  'They will be removed in JAX 0.22.2', FutureWarning)
+  if min is None and a_min is not None:
+    min = a_min
+  if max is None and a_max is not None:
+    max = a_max
+  return clip(number, a_min=min, a_max=max, out=out)
+
+
 def _view(arr, dtype=None, type=None):
   lax._check_user_dtype_supported(dtype, "view")
   if type is not None:
@@ -6082,7 +6096,7 @@ _operators = {
 # These numpy.ndarray methods are just refs to an equivalent numpy function
 _nondiff_methods = ["all", "any", "argmax", "argmin", "argpartition", "argsort",
                     "nonzero", "searchsorted", "round"]
-_diff_methods = ["choose", "clip", "conj", "conjugate", "cumprod", "cumsum",
+_diff_methods = ["choose", "conj", "conjugate", "cumprod", "cumsum",
                  "diagonal", "dot", "max", "mean", "min", "prod", "ptp",
                  "ravel", "repeat", "sort", "squeeze", "std", "sum",
                  "swapaxes", "take", "tile", "trace", "var"]
@@ -6313,6 +6327,7 @@ def _set_shaped_array_attributes(shaped_array):
   setattr(shaped_array, "astype", core.aval_method(_astype))
   setattr(shaped_array, "view", core.aval_method(_view))
   setattr(shaped_array, "nbytes", core.aval_property(_nbytes))
+  setattr(shaped_array, "clip", core.aval_method(_clip))
 
   setattr(shaped_array, "_array_module", staticmethod(__array_module__))
   setattr(shaped_array, "broadcast", core.aval_method(lax.broadcast))
@@ -6340,6 +6355,7 @@ def _set_device_array_base_attributes(device_array):
   setattr(device_array, "astype", _astype)
   setattr(device_array, "view", _view)
   setattr(device_array, "nbytes", property(_nbytes))
+  setattr(device_array, "clip", _clip)
 
 _set_device_array_base_attributes(DeviceArray)
Repository: google/jax
Issue: Warning when using jnp.any, jnp.all on complex arrays

The `numpy` equivalents operate quietly on complex arrays, while the `jax.numpy` variants raise a ComplexWarning

```python 
In [1]: x = np.random.randn(2) + 1j * np.random.randn(2)

In [2]: np.any(x)
Out[2]: True

In [3]: jax.numpy.any(jax.device_put(x))
/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py:1945: ComplexWarning: Casting complex values to real discards the imaginary part
  a = preproc(a) if preproc else a
Out[3]: DeviceArray(True, dtype=bool)
```
- jax 0.2.11
- jaxlib 0.1.64

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 95148f016..7e3306743 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1993,7 +1993,10 @@ def _reduction_init_val(a, init_val):
     sign, info = np.sign(init_val), iinfo(a_dtype)
     return np.array(info.min if sign < 0 else info.max, dtype=a_dtype)
 
-_cast_to_bool = partial(lax.convert_element_type, new_dtype=bool_)
+def _cast_to_bool(operand):
+  with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=np.ComplexWarning)
+    return lax.convert_element_type(operand, bool_)
 
 @_wraps(np.sum, skip_params=['out'])
 def sum(a, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype=None,
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 501ab9e1c..36d6725dc 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -785,9 +785,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       res = res if not is_bf16_nan_test else res.astype(jnp.bfloat16)
       return res
     np_fun = _promote_like_jnp(np_fun, inexact)
-    np_fun = jtu.ignore_warning(category=np.ComplexWarning)(np_fun)
     jnp_fun = lambda x: jnp_op(x, axis, keepdims=keepdims)
-    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
     args_maker = lambda: [rng(shape, dtype)]
     tol = {np.float16: 0.002}
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, tol=tol)
Repository: google/jax
Issue: jnp.flip crashes when given a tuple of axes

Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

```python
import jax.numpy as jnp
jnp.flip(jnp.zeros((3, 3, 3)), axis=(1,2))
```

- [x] If applicable, include full error messages/tracebacks.
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/buehrle/dev/polyssimo/.venv/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py", line 1114, in flip
    return lax.rev(m, [_canonicalize_axis(axis, ndim(m))])
  File "/home/buehrle/dev/polyssimo/.venv/lib/python3.6/site-packages/jax/_src/util.py", line 262, in canonicalize_axis
    axis = operator.index(axis)
TypeError: 'tuple' object cannot be interpreted as an integer
```

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 582eb4f73..ae6533530 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1111,7 +1111,8 @@ def flip(m, axis: Optional[Union[int, Tuple[int, ...]]] = None):
   _check_arraylike("flip", m)
   if axis is None:
     return lax.rev(m, list(range(len(shape(m)))))
-  return lax.rev(m, [_canonicalize_axis(axis, ndim(m))])
+  axis = _ensure_index_tuple(axis)
+  return lax.rev(m, [_canonicalize_axis(ax, ndim(m)) for ax in axis])
 
 
 @_wraps(np.fliplr)
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 422719e3b..2b540114b 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -3449,7 +3449,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
        "shape": shape, "dtype": dtype, "axis": axis}
       for shape in [(3,), (2, 3)]
       for dtype in default_dtypes
-      for axis in list(range(-len(shape), len(shape))) + [None]  # Test negative axes
+      for axis in list(range(-len(shape), len(shape))) + [None] + [tuple(range(len(shape)))]  # Test negative axes and tuples
     ))
   def testFlip(self, shape, dtype, axis):
     rng = jtu.rand_default(self.rng())
Repository: google/jax
Issue: Gradient through Neural Tangent Kernel works for CNN without MaxPool, but not for CNN with MaxPool

Hi all,

I am working on a project that requires taking the gradient of Neural Tangent Kernel with respect to the parameters at which the Jacobian is being evaluated.

However, I’m getting an error when jitting the NTK computation function if my network uses maxpooling.

I have prepared two Colab notebooks that reproduce the issue:
[CNN without MaxPool](https://colab.research.google.com/drive/1k_cCuNxz2rXeXFrVnwFimWtvOzfVPjdo?usp=sharing)
[CNN with MaxPool](https://colab.research.google.com/drive/1WxK6WjSrZ4FdzJpY3hg0LSZ8gSkM7Bf-?usp=sharing)

In the first notebook, I show that computing the gradients with respect to the NTK evaluation parameters works just fine if I don’t include a maxpooling operation in the CNN, and in the second notebook, I show that including a maxpooling operation results in an error:
> typeValue Zero(ShapedArray(float32[11,28,28,32])) with type <class ‘jax.ad_util.Zero’> is not a valid JAX type

Oddly, the error message disappears if I remove all the jitting from the NTK function.

I would really appreciate if you could help me understand the cause of this issue and help me fix it.

Thank you so much!

Changed files:
jax/_src/lax/lax.py

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index f8944b13a..e2df6b7b0 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -5573,6 +5573,8 @@ def _select_and_scatter_add_transpose(
     t, source, operand, *, select_prim, window_dimensions, window_strides,
     padding):
   assert ad.is_undefined_primal(source) and not ad.is_undefined_primal(operand)
+  if type(t) is ad_util.Zero:
+    return [ad_util.Zero(source.aval), None]
   ones = (1,) * len(window_dimensions)
   source_t = _select_and_gather_add(t, operand, select_prim, window_dimensions,
                                     window_strides, padding, ones, ones)
Repository: google/jax
Issue: Inconsistent edge case handling in jax.random.categorical across devices

When provided NaN probabilities, `jax.random.categorical` returns different results depending on the device. 

```
# GPU
> jax.random.categorical(rng_key, jnp.array([0, jnp.nan]))
DeviceArray(2147483647, dtype=int32)

# CPU
> jax.random.categorical(rng_key, jnp.array([0, jnp.nan]))
DeviceArray(1, dtype=int32)
```

Arguably both are not ideal (being valid values, bugs can go unnoticed - especially in the CPU case where 1 is a perfectly valid value coming from this function) but the inconsistency in particular can make debugging very tricky and should be fixed.

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

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 05fb817f7..8e1ff56f7 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -5382,7 +5382,8 @@ def _argminmax_translation_rule(value_comparator, identity,
   x_index = xb.parameter(subc, 1, index_shape)
   y_value = xb.parameter(subc, 2, value_shape)
   y_index = xb.parameter(subc, 3, index_shape)
-  which_value = value_comparator(x_value, y_value)
+  which_value = xops.Or(value_comparator(x_value, y_value),
+                        xops.Ne(x_value, x_value))
   which_index = xops.Or(which_value, xops.And(xops.Eq(x_value, y_value),
                                               xops.Lt(x_index, y_index)))
   xops.Tuple(subc, [xops.Select(which_value, x_value, y_value),
@@ -5402,8 +5403,8 @@ def _argminmax_gpu_translation_rule(op, a, *, axes, index_dtype):
   idxs = tie_in(a, broadcasted_iota(index_dtype, a.shape, axis))
   maxval = np.array(dtypes.iinfo(index_dtype).max, dtype=index_dtype)
   maxval = broadcast(tie_in(a, maxval), a.shape)
-  mask_idxs = select(eq(a, expand_dims(op(a, (axis,)), (axis,))), idxs,
-                     maxval)
+  maxvals = expand_dims(op(a, (axis,)), (axis,))
+  mask_idxs = select(eq(a, maxvals) | ne(a, a), idxs, maxval)
   return _reduce_min(mask_idxs, (axis,))
 
 _argmin_translation_rule = partial(_argminmax_translation_rule, xops.Lt,
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 9ca26b04e..15817ae75 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2560,6 +2560,10 @@ class LazyConstantTest(jtu.JaxTestCase):
     x_out_jit = api.jit(op)(x_in)
     self.assertEqual(dtypes.is_weakly_typed(x_out_jit), False)
 
+  def testArgMaxOfNanChoosesNaN(self):
+    self.assertEqual(lax.argmax(np.array([0., np.nan]), axis=0,
+                                index_dtype=np.int32), 1)
+
   @parameterized.named_parameters(jtu.cases_from_list(
         {"testcase_name": "_{}".format(rec.op),
          "op_name": rec.op, "rec_dtypes": rec.dtypes}
Repository: google/jax
Issue: tuple_replace function in util.py seems incorrect

In the file jax._src.util.py

Current definition
```python
def tuple_replace(t, idx, val):
  assert 0 <= idx < len(t), (idx, len(t))
  return t[:idx] + (val,) + t[idx:]
```
looks like an almost copy of tuple_insert. I think the correct definition should be:

```python
def tuple_replace(t, idx, val):
  assert 0 <= idx < len(t), (idx, len(t))
  return t[:idx] + (val,) + t[idx+1:]
```



Changed files:
jax/_src/util.py

diff --git a/jax/_src/util.py b/jax/_src/util.py
index 607479216..051f570f3 100644
--- a/jax/_src/util.py
+++ b/jax/_src/util.py
@@ -330,10 +330,6 @@ def tuple_delete(t, idx):
   assert 0 <= idx < len(t), (idx, len(t))
   return t[:idx] + t[idx + 1:]
 
-def tuple_replace(t, idx, val):
-  assert 0 <= idx < len(t), (idx, len(t))
-  return t[:idx] + (val,) + t[idx:]
-
 # TODO(mattjj): replace with dataclass when Python 2 support is removed
 def taggedtuple(name, fields) -> Callable[..., Any]:
   """Lightweight version of namedtuple where equality depends on the type."""
Repository: google/jax
Issue: Support negative numpy.split indices

```python
import numpy as np
np.split([1, 2, 3], [-1])
```

works, while

```python
import jax.numpy as jnp
jnp.split(jnp.array([1, 2, 3]), [-1])
```

fails.

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 543da9681..fa29438d7 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1805,7 +1805,8 @@ def _split(op, ary, indices_or_sections, axis=0):
   size = ary.shape[axis]
   if isinstance(indices_or_sections, (tuple, list) + _arraylike_types):
     indices_or_sections = np.array(
-        [core.concrete_or_error(np.int64, i_s, f"in jax.numpy.{op} argument 1")
+        [_canonicalize_axis(core.concrete_or_error(np.int64, i_s,
+                                                   f"in jax.numpy.{op} argument 1"), size)
          for i_s in indices_or_sections], np.int64)
     split_indices = np.concatenate([[np.int64(0)], indices_or_sections,
                                     [np.int64(size)]])
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 6e2d1799d..b75cde9cf 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2925,8 +2925,8 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
        "shape": shape, "num_sections": num_sections, "axis": axis,
        "dtype": dtype}
       for shape, axis, num_sections in [
-          ((3,), 0, 3), ((12,), 0, 3), ((12, 4), 0, 4), ((12, 4), 1, 2),
-          ((2, 3, 4), -1, 2), ((2, 3, 4), -2, 3)]
+          ((3,), 0, 3), ((12,), 0, 3), ((12, 4), 0, 4), ((3,), 0, [-1]),
+          ((12, 4), 1, 2), ((2, 3, 4), -1, 2), ((2, 3, 4), -2, 3)]
       for dtype in default_dtypes))
   def testSplitStaticInt(self, shape, num_sections, axis, dtype):
     rng = jtu.rand_default(self.rng())