anon-iclr-submission/benchname-bug-localization huggingface dataset:
Repository: google/jax
Issue: tree_util deprecation warnings should set warning class and stacklevel

If I try using deprecated functions like `jax.tree_map` using the development version of JAX, I see messages like:
```
>>> jax.tree_map(lambda x: x, [1, 2, 3])
/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py:529: UserWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
  warnings.warn(f"jax.{f.__name__} is deprecated, and will be removed in a future release. "
[1, 2, 3]
```

This is OK, but a better warning would set `category` (so it can be filtered appropriately) and `stacklevel` (so it's clear to users where it's coming from).

Either `DeprecationWarning` or `FutureWarning` would be appropriate for the category, depending on how "noisy" you want the warnings to be. I would probably lean towards `FutureWarning` for JAX, since JAX is typically used by developers rather than end users.

Changed files:
jax/_src/tree_util.py

diff --git a/jax/_src/tree_util.py b/jax/_src/tree_util.py
index 344edca8e..353af2a12 100644
--- a/jax/_src/tree_util.py
+++ b/jax/_src/tree_util.py
@@ -527,7 +527,8 @@ def _deprecate(f):
   @functools.wraps(f)
   def wrapped(*args, **kwargs):
     warnings.warn(f"jax.{f.__name__} is deprecated, and will be removed in a future release. "
-                  f"Use jax.tree_util.{f.__name__} instead.")
+                  f"Use jax.tree_util.{f.__name__} instead.",
+                  category=FutureWarning, stacklevel=2)
     return f(*args, **kwargs)
   return wrapped
Repository: google/jax
Issue: Misleading error message in broadcast

### Description

```python
In [64]: x, y = jnp.ones((2, 4)), jnp.ones((2,))
In [65]: x + y
....
/jax/_src/numpy/util.py in _promote_shapes(fun_name, *args)
    246         if config.jax_numpy_rank_promotion != "allow":
    247           _rank_promotion_warning_or_error(fun_name, shapes)
--> 248         result_rank = len(lax.broadcast_shapes(*shapes))
    249         return [_broadcast_to(arg, (1,) * (result_rank - len(shp)) + shp)
    250                 for arg, shp in zip(args, shapes)]

ValueError: Incompatible shapes for broadcasting: ((2, 4), (1, 2))
```
```python
In [66]: xx, yy = np.ones((2, 4)), np.ones((2,))
In [67]: xx + yy
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-67-9e546835689c> in <module>()
----> 1 xx + yy

ValueError: operands could not be broadcast together with shapes (2,4) (2,)
```

The error message from `jnp` does not tell the shape given by user and that's misleading for debugging. `np` has a better behavior here.

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

latest

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

_No response_

### Additional System Info

_No response_

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

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index e709eaa19..5bc6c504a 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -149,8 +149,7 @@ def _broadcast_shapes_uncached(*shapes):
   shape_list = [(1,) * (ndim - len(shape)) + shape for shape in shapes]
   result_shape = _try_broadcast_shapes(shape_list)
   if result_shape is None:
-    raise ValueError("Incompatible shapes for broadcasting: {}"
-                     .format(tuple(shape_list)))
+    raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
   return result_shape
 
 def _broadcast_ranks(s1, s2):
Repository: google/jax
Issue: improve error message for jnp.diagonal

```python
import jax.numpy as jnp
import numpy as np
M = jnp.ones(5)
```

Then, `np.diagonal(M)` returns a clear error message `ValueError: diag requires an array of at least two dimensions` whereas `jnp.diagonal(M)` returns a still-informative-but-less-clear `ValueError: axis 1 is out of bounds for array of dimension 1`.

Could we possibly change this behavior?

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 1c63e3553..f477e4c17 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2516,6 +2516,8 @@ def diag_indices_from(arr):
 def diagonal(a, offset=0, axis1: int = 0, axis2: int = 1):
   _check_arraylike("diagonal", a)
   a_shape = shape(a)
+  if ndim(a) < 2:
+    raise ValueError("diagonal requires an array of at least two dimensions.")
   offset = core.concrete_or_error(operator.index, offset, "'offset' argument of jnp.diagonal()")
 
   a = moveaxis(a, (axis1, axis2), (-2, -1))
Repository: google/jax
Issue: `lax.rem` has wrong gradient for `y` when `x<0`

`lax.rem` has wrong gradient when `x<0`. For example,

```py
import jax

def fn(y):
    x = jax.numpy.array(-19., jax.numpy.float64)
    return jax.lax.rem(x, y)

y = jax.numpy.array(3., jax.numpy.float64)
res = fn(y)
print(res) # -1.0

a = jax.jacfwd(fn)(y)
print(a) # 7.0
```

Actually, the gradient should be 6 instead of 7. Because `-19 = (-6)*3 + (-1)`. When the `y` is `3.1`, the results will be `-0.4`. Thus, the gradient should be 6.

By the way, another two similar APIs `numpy.remainder` and `numpy.mod` also have this issue.

```py
def fn(y):
    x = jax.numpy.array(-19., jax.numpy.float64)
    return jax.numpy.remainder(x, y)

y = jax.numpy.array(3., jax.numpy.float64)
res = fn(y)
print(res) # 2.0

a = jax.jacfwd(fn)(y)
print(a) # 8.0
```

In these two APIs, `-19 = (-7)*3 + 2`, so their gradient should be 7 instead of 8.

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

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index aeb80847c..3980d79bc 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2161,7 +2161,7 @@ rem_p = standard_naryop([_int | _float, _int | _float], 'rem')
 ad.defjvp(
     rem_p,
     lambda g, x, y: _maybe_broadcast(broadcast_shapes(np.shape(x), np.shape(y)), g),
-    lambda g, x, y: mul(neg(g), floor(div(x, y))))
+    lambda g, x, y: mul(neg(g), mul(sign(div(x, y)), floor(abs(div(x, y))))))
 mlir.register_lowering(rem_p, partial(_nary_lower_mhlo, mhlo.RemOp))
 
 def _minmax_complex_lowering(x, y, *, lax_cmp_pick_x):
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
index 155fab36e..41346ae45 100644
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -1060,19 +1060,24 @@ class LaxAutodiffTest(jtu.JaxTestCase):
 
   # TODO(mattjj): make this a more systematic test
   def testRemainder(self):
+    def gen_x(rng, size):
+      return rng.uniform(-9, 9, size=size)
+
+    def gen_y(rng, size):
+      # avoid values near zero because gradients diverge
+      return rng.uniform(0.1, 5, size=size) * rng.choice([-1, 1], size=size)
+
     rng = self.rng()
-    x = rng.uniform(-0.9, 9, size=(3, 4))
-    y = rng.uniform(0.7, 1.9, size=(3, 1))
+    x = gen_x(rng, (5, 8))
+    y = gen_y(rng, (1, 8))
     assert not set(np.unique(x)) & set(np.unique(y))
-    # TODO(jakevdp) try to make these tolerances tighter.
-    tol = 1e-1
-    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"], tol, tol)
+    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"])
 
     rng = self.rng()
-    x = rng.uniform(-0.9, 9, size=(1, 4))
-    y = rng.uniform(0.7, 1.9, size=(3, 4))
+    x = gen_x(rng, (1, 8))
+    y = gen_y(rng, (5, 8))
     assert not set(np.unique(x)) & set(np.unique(y))
-    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"], tol, tol)
+    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"])
 
   def testHigherOrderGradientOfReciprocal(self):
     # Regression test for https://github.com/google/jax/issues/3136
Repository: google/jax
Issue: `lax.conv_general_dilated` will crash when `padding` has some negative

`lax.conv_general_dilated` will crash when `padding` has some negative

```py
import jax
def fn(lhs, rhs):
    window_strides = [1, 1]
    padding = [[5, 16], [-5, 1]]
    return jax.lax.conv_general_dilated(lhs, rhs, window_strides, padding)

mykey = jax.random.PRNGKey(83725540)

lhs = jax.numpy.array([[[[-7.0950737, 20.612083, 53.864067], [-12.116554, -16.742317, -7.050851 ]]]])
rhs = jax.numpy.array([[[[29.074692, 40.518192], [21.685156, 41.867775]]]])

fn(lhs, rhs)
```
```
F external/org_tensorflow/tensorflow/compiler/xla/window_util.cc:250] Check failed: bound >= 0 (0 vs. -1)
[1]   abort (core dumped)
```

Besides, `jax.lax.conv_general_dilated_patches` and `jax.lax.conv_with_general_padding` also have such issue

```py
import jax
def fn(lhs):
    filter_shape = [1, 1]
    window_strides = [1, 1]
    padding = [[-4, -2], [-3, -4]]
    return jax.lax.conv_general_dilated_patches(lhs=lhs, filter_shape=filter_shape, window_strides=window_strides, padding=padding)

mykey = jax.random.PRNGKey(19446744)
lhs = jax.random.uniform(mykey, [1, 2, 3, 4], jax.numpy.float32)
fn(lhs)
```

```
F external/org_tensorflow/tensorflow/compiler/xla/window_util.cc:250] Check failed: bound >= 0 (0 vs. -3)
[1]   abort (core dumped)
```

```py
import jax
def fn(lhs, rhs):
    window_strides = [1, 1]
    padding = [[-5, 0], [-16, -2]]
    lhs_dilation = [1, 1]
    rhs_dilation = [1, 1]
    return jax.lax.conv_with_general_padding(lhs, rhs, window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, rhs_dilation=rhs_dilation)

mykey = jax.random.PRNGKey(30297051)
lhs = jax.random.uniform(mykey, [3, 2, 3, 1], jax.numpy.float32)
mykey = jax.random.PRNGKey(34585797)
rhs = jax.random.uniform(mykey, [2, 2, 1, 2], jax.numpy.float32)

fn(lhs, rhs)
```

```
F external/org_tensorflow/tensorflow/compiler/xla/window_util.cc:250] Check failed: bound >= 0 (0 vs. -2)
[1]   abort (core dumped)
```

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

diff --git a/jax/_src/lax/convolution.py b/jax/_src/lax/convolution.py
index e73d147ad..1c31fe6a3 100644
--- a/jax/_src/lax/convolution.py
+++ b/jax/_src/lax/convolution.py
@@ -812,6 +812,9 @@ def conv_shape_tuple(lhs_shape, rhs_shape, strides, pads, batch_group_count=1):
 
   lhs_padded = np.add(lhs_shape[2:], np.sum(np.array(pads).reshape(-1, 2),
                                               axis=1))
+  if np.any(lhs_padded < 0):
+    raise ValueError("Negative padding is larger than the size of the corresponding dimension: "
+                     f"got padding={pads} for lhs_shape[2:]={lhs_shape[2:]}")
   out_space = core.stride_shape(lhs_padded, rhs_shape[2:], strides)
   out_space = np.maximum(0, out_space)
   if batch_group_count > 1:
Repository: google/jax
Issue: BUG: numpy and jax differ on modulus by zero

### Description

Jax and numpy differ in their behavior with regards to computing a modulus by zero. Jax treats it as a no-op, and numpy treats it as the zero function.

## Numpy

```python
Python 3.10.5 (main, Jul 15 2022, 03:56:49) [GCC 11.3.0] on linux
>>> import numpy as np
>>> a = np.array([1,2,3], dtype=np.uint32)
>>> np.mod(a, 0)
<stdin>:1: RuntimeWarning: divide by zero encountered in remainder
array([0, 0, 0], dtype=uint32)
```

## JAX

```python
>>> import jax.numpy as jnp
>>> a = jnp.array([1,2,3], dtype=jnp.uint32)
>>> jnp.mod(a, 0)
DeviceArray([1, 2, 3], dtype=uint32)
```

### What jax/jaxlib version are you using?

HEAD

### Which accelerator(s) are you using?

CPU

### Additional System Info

Python 3.9, 3.10.5

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

diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index 1062650e3..5795a0fda 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -529,6 +529,8 @@ def frexp(x):
 def remainder(x1, x2):
   x1, x2 = _promote_args_numeric("remainder", x1, x2)
   zero = _constant_like(x1, 0)
+  if dtypes.issubdtype(x2.dtype, np.integer):
+    x2 = _where(x2 == 0, lax_internal._ones(x2), x2)
   trunc_mod = lax.rem(x1, x2)
   trunc_mod_not_zero = lax.ne(trunc_mod, zero)
   do_plus = lax.bitwise_and(
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index d371eedec..b6362783c 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -286,9 +286,9 @@ JAX_COMPOUND_OP_RECORDS = [
     op_record("rad2deg", 1, float_dtypes, all_shapes, jtu.rand_default, []),
     op_record("ravel", 1, all_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("real", 1, number_dtypes, all_shapes, jtu.rand_some_inf, []),
-    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero, [],
+    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_some_zero, [],
               tolerance={np.float16: 1e-2}),
-    op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_nonzero, []),
+    op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_some_zero, []),
     op_record("modf", 1, float_dtypes, all_shapes, jtu.rand_default, []),
     op_record("modf", 1, int_dtypes + unsigned_dtypes, all_shapes,
               jtu.rand_default, [], check_dtypes=False),
Repository: google/jax
Issue: jax.lax.cond evaluates None as True

When using `jax.lax.cond`, passing `None` as the first argument causes the `true_fun` argument to get called. This is contrary to the behavior one would expect for Python if statements.


For example, running the following
```python
import jax

true_fun = lambda x: True
false_fun = lambda x: False

for pred in [True, False, None]:
	print(pred, jax.lax.cond(pred, true_fun, false_fun, None))
```

gives the following output:
```
True True
False False
None True
```

I am using the newest PyPI versions of `jax` and `jaxlib` (0.3.14).

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

diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
index d88d82700..b94ed59d4 100644
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -148,7 +148,8 @@ def _cond(pred, true_fun: Callable, false_fun: Callable, *operands,
           operand=_no_operand_sentinel, linear=None):
   """Conditionally apply ``true_fun`` or ``false_fun``.
 
-  ``cond()`` has equivalent semantics to this Python implementation::
+  Provided arguments are correctly typed, ``cond()`` has equivalent
+  semantics to this Python implementation::
 
     def cond(pred, true_fun, false_fun, *operands):
       if pred:
@@ -181,6 +182,8 @@ def _cond(pred, true_fun: Callable, false_fun: Callable, *operands,
     operands = (operand,)
   del operand
 
+  if pred is None:
+    raise TypeError("cond predicate is None")
   if isinstance(pred, Sequence) or np.ndim(pred) != 0:
     raise TypeError(
         f"Pred must be a scalar, got {pred} of " +
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index 6fddbd437..ce07f65ad 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -628,6 +628,16 @@ class LaxControlFlowTest(jtu.JaxTestCase):
     self.assertEqual(fun(4), cfun(4))
     self.assertEqual(fun(4), (8, 16))
 
+  def testCondPredIsNone(self):
+    # see https://github.com/google/jax/issues/11574
+    def f(pred, x):
+      return lax.cond(pred, lambda x: x + 1, lambda x: x + 2, x)
+
+    self.assertRaisesRegex(TypeError, "cond predicate is None",
+                           lambda: f(None, 1.))
+    self.assertRaisesRegex(TypeError, "cond predicate is None",
+                           lambda: jax.jit(f)(None, 1.))
+
   def testCondTwoOperands(self):
     # see https://github.com/google/jax/issues/8469
     add, mul = lax.add, lax.mul
Repository: google/jax
Issue: BUG: jnp.linalg.matrix_rank does not operate on stacks of matrices

### Description

```
myInput = np.random.normal(size=(9,10,10))
myInput.shape
# (9,10,10)
np.linalg.matrix_rank(myInput)
# array([10, 10, 10, 10, 10, 10, 10, 10, 10])
```

jax fails on this same stack:
```
jax.numpy.linalg.matrix_rank(jnp.array(myInput))
### TypeError: array should have 2 or fewer dimensions
### Expected: DeviceArray([10, 10, 10, 10, 10, 10, 10, 10, 10])
```

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

jax 0.3.13, jaxlib 0.3.10

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

- [ ] CPU
- [X] GPU
- [ ] TPU

### Additional System Info

_No response_

Changed files:
tests/linalg_test.py, jax/_src/numpy/linalg.py

diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index 3f3b172bc..97827576e 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -111,14 +111,13 @@ def matrix_power(a, n):
 @jit
 def matrix_rank(M, tol=None):
   M, = _promote_dtypes_inexact(jnp.asarray(M))
-  if M.ndim > 2:
-    raise TypeError("array should have 2 or fewer dimensions")
   if M.ndim < 2:
     return jnp.any(M != 0).astype(jnp.int32)
   S = svd(M, full_matrices=False, compute_uv=False)
   if tol is None:
-    tol = S.max() * np.max(M.shape).astype(S.dtype) * jnp.finfo(S.dtype).eps
-  return jnp.sum(S > tol)
+    tol = S.max(-1) * np.max(M.shape[-2:]).astype(S.dtype) * jnp.finfo(S.dtype).eps
+  tol = jnp.expand_dims(tol, np.ndim(tol))
+  return jnp.sum(S > tol, axis=-1)
 
 
 @custom_jvp
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 97234906f..ff9f28c95 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -924,7 +924,8 @@ class NumpyLinalgTest(jtu.JaxTestCase):
       {"testcase_name": "_shape={}".format(
            jtu.format_shape_dtype_string(shape, dtype)),
        "shape": shape, "dtype": dtype}
-      for shape in [(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50)]
+      for shape in [(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50),
+                    (3, 4, 5), (2, 3, 4, 5)]
       for dtype in float_types + complex_types))
   def testMatrixRank(self, shape, dtype):
     rng = jtu.rand_default(self.rng())
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
```


Changed files:
jax/_src/nn/functions.py, tests/nn_test.py

diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
index 3d568983b..bc386bb7c 100644
--- 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.
 
diff --git a/tests/nn_test.py b/tests/nn_test.py
index 1659f7d40..60b0764a8 100644
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -81,7 +81,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
     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: [bug] using `lax.cond` with partial effect function gives TypeError

In the master branch, the following script shows the error
```sh
TypeError: could not merge true_fun and false_fun output pvals: ((ShapedArray(int32[]), None), JaxTuple(*,0)) and ((ShapedArray(int32[]),ShapedArray(int32[])), JaxTuple(*,*)).
```
```
lax.cond(True, (0, 0), lambda x: (x[0], 0), (1, 1), lambda x: x)

```

While these versions works well:
```
lax.cond(True, (0, 0), lambda x: (0, 0), (1, 1), lambda x: x)
```
or
```
lax.cond(True, (0, 0), lambda x: (x[0], x[1]), (1, 1), lambda x: x)
```

Changed files:
jax/interpreters/partial_eval.py, tests/lax_test.py

diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 91bdc883a..f49966a6d 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -296,7 +296,7 @@ def join_pvals(pval1, pval2):
     pvals1, pvals2 = zip(pv1, const1), zip(pv2, const2)
     join_pvs, join_consts = unzip2(map(join_pvals, pvals1, pvals2))
     if all(isinstance(pv, AbstractValue) for pv in join_pvs):
-      return PartialVal(AbstractTuple(join_pvs), tuple(join_consts))
+      return PartialVal((AbstractTuple(join_pvs), tuple(join_consts)))
     else:
       return PartialVal((JaxprTracerTuple(join_pvs), tuple(join_consts)))
 
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 7c095c6c8..45039f759 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1421,6 +1421,12 @@ class LaxTest(jtu.JaxTestCase):
     self.assertEqual(fun(4), cfun(4))
     self.assertEqual(cfun(4), (4, 2., 4.))
 
+  def testIssue514(self):
+    # just check this doesn't crash
+    lax.cond(True,
+            (0, 0), lambda x: (x[0], 0),
+            (1, 1), lambda x: x)
+
   def testScanAdd(self):
     def f(x, y):
       return x + y
Repository: google/jax
Issue: AttributeError: module 'jaxlib.xla_extension' has no attribute 'PmapFunction'

I am getting the above error after installing the latest version of jax and trying to import it for some computations:
```
!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
```
```
/usr/local/lib/python3.7/dist-packages/jax/_src/lax/linalg.py in <module>()
    977       lu_pivots_to_permutation_p,
    978       partial(_lu_pivots_to_permutation_gpu_lowering,
--> 979               gpu_linalg.cuda_lu_pivots_to_permutation),
    980       platform='cuda')
    981   mlir.register_lowering(

AttributeError: module 'jaxlib.cuda_linalg' has no attribute 'cuda_lu_pivots_to_permutation'
```

It would be helpful if someone could please look into this ASAP. This is a new error I received when running my code within the past hour and had no problem with this earlier. 

Changed files:
jax/_src/lib/__init__.py

diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
index ea9b2d9a8..f45bf7295 100644
--- a/jax/_src/lib/__init__.py
+++ b/jax/_src/lib/__init__.py
@@ -183,7 +183,7 @@ except ImportError:
   hip_linalg = None
 
 try:
-  import jaxlib.cuda_linalg as gpu_linalg  # pytype: disable=import-error
+  import jaxlib.gpu_linalg as gpu_linalg  # pytype: disable=import-error
 except ImportError:
   gpu_linalg = None
Repository: google/jax
Issue: undefined f if vmap applied to lu.WrappedFun

In `vmap`, `f` is not be defined whenever the if condition is false:

https://github.com/google/jax/blob/f8eda9bb645a1ed489f08be57bae380bbcf5fa0c/jax/api.py#L414-L420

Changed files:
jax/api.py

diff --git a/jax/api.py b/jax/api.py
index 03f925880..ed4652ed6 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -413,8 +413,7 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   @wraps(fun, docstr=docstr)
   def batched_fun(*args, **kwargs):
-    if not isinstance(fun, lu.WrappedFun):
-      f = lu.wrap_init(fun, kwargs)
+    f = lu.wrap_init(fun, kwargs) if not isinstance(fun, lu.WrappedFun) else fun
     in_axes_ = in_axes if isinstance(in_axes, (list, tuple)) else (in_axes,) * len(args)
     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
Repository: google/jax
Issue: generators should not be accepted as values for the shape argument of numpy functions

```python3
import numpy as np
shape = (1 for _ in range(4)) 
np.zeros(shape)
```

results in `TypeError: expected sequence object with len >= 0 or a single integer`

whereas

```python3
import jax.numpy as np
shape = (1 for _ in range(4)) 
np.zeros(shape)
```

just runs through and produces `array([[[[0.]]]], dtype=float32)`.
This however can lead to bugs that are hard to find, because calling `np.zeros(shape)` a second time will create `array(0., dtype=float32)`

Conclusion: `jax.numpy` should behave like `numpy` and should also through a TypeError

Changed files:
jax/numpy/lax_numpy.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index de22537ce..20f7be4a9 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -21,6 +21,7 @@ import itertools
 import re
 import string
 import warnings
+import types
 
 import numpy as onp
 import opt_einsum
@@ -1221,6 +1222,8 @@ def full_like(a, fill_value, dtype=None):
 
 @_wraps(onp.zeros)
 def zeros(shape, dtype=onp.dtype("float64")):
+  if isinstance(shape, types.GeneratorType):
+    raise TypeError("expected sequence object with len >= 0 or a single integer")
   shape = (shape,) if onp.isscalar(shape) else shape
   return lax.full(shape, 0, dtype)
Repository: google/jax
Issue: Regression in performance of samplers

With jax versions > 0.1.24, random samplers are slow. I think that it is due to recent changes on how jit works with static_args/kwargs. A script to reproduce,
```
import time
from jax import random

t = time.time()
random.normal(random.PRNGKey(0), shape=(1,))
print(time.time() - t)
t = time.time()
random.normal(random.PRNGKey(0), shape=(1,))
print(time.time() - t)
```
which returns `0.12923526763916016` and `0.12221717834472656`.

However, if we wrap these samplers in some function, then it is fast. For example,
```
def f():
    return random.normal(random.PRNGKey(0), shape=(1,))

t = time.time()
f()
print(time.time() - t)
t = time.time()
f()
print(time.time() - t)
```
will return `0.12787413597106934` and `0.0010831356048583984`.

I think that there is a small bug elsewhere which forces the sampler recompile. If this is an expected behaviour, then which function we should use to wrap these samplers to make it not recompiled?
cc @neerajprad


Changed files:
jax/util.py, jax/api.py

diff --git a/jax/api.py b/jax/api.py
index 7d7ae951d..5ce6f909d 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -46,7 +46,7 @@ from .tree_util import (process_pytree, node_types, build_tree, PyTreeDef,
                         tree_map, tree_flatten, tree_unflatten, tree_structure,
                         tree_transpose, leaf)
 from .util import (unzip2, unzip3, curry, partial, safe_map, safe_zip,
-                   WrapHashably, prod)
+                   WrapHashably, Hashable, prod)
 from .lib.xla_bridge import canonicalize_dtype, device_count
 from .abstract_arrays import ShapedArray
 from .interpreters import partial_eval as pe
@@ -824,11 +824,19 @@ def _argnums_partial(f, dyn_argnums, args):
     dyn_argnums = (dyn_argnums,)
   else:
     dyn_argnums = tuple(dyn_argnums)
-  fixed_args = tuple([None if i in dyn_argnums else WrapHashably(arg)
+  fixed_args = tuple([None if i in dyn_argnums else _wrap_hashably(arg)
                       for i, arg in enumerate(args)])
   dyn_args = tuple(args[i] for i in dyn_argnums)
   return _argnums_partial_(f, dyn_argnums, fixed_args), dyn_args
 
+def _wrap_hashably(arg):
+  try:
+    hash(arg)
+  except TypeError:
+    return WrapHashably(arg)
+  else:
+    return Hashable(arg)
+
 @lu.transformation
 def _argnums_partial_(dyn_argnums, fixed_args, *dyn_args, **kwargs):
   args = [None if arg is None else arg.val for arg in fixed_args]
diff --git a/jax/util.py b/jax/util.py
index 8faea1361..f54cb022d 100644
--- a/jax/util.py
+++ b/jax/util.py
@@ -189,6 +189,8 @@ def prod(xs):
 
 
 class WrapHashably(object):
+  __slots__ = ["val"]
+
   def __init__(self, val):
     self.val = val
 
@@ -198,6 +200,18 @@ class WrapHashably(object):
   def __eq__(self, other):
     return self.val is other.val
 
+class Hashable(object):
+  __slots__ = ["val"]
+
+  def __init__(self, val):
+    self.val = val
+
+  def __hash__(self):
+    return hash(self.val)
+
+  def __eq__(self, other):
+    return self.val == other.val
+
 
 
 def get_module_functions(module):
Repository: google/jax
Issue: device_put leaks references on numpy inputs

First reported in #14876; here's a minimal repro that reproduces on CPU and GPU:
```python
import sys
import jax
import numpy as np

x = np.arange(1000, dtype='float32')
print(sys.getrefcount(x))
# 2

y = jax.device_put(x)
print(sys.getrefcount(x))
# 3

del y
print(sys.getrefcount(x))
# 3
```

Changed files:
jax/_src/lib/__init__.py, tests/api_test.py

diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
index 6c434d53c..3144402c8 100644
--- a/jax/_src/lib/__init__.py
+++ b/jax/_src/lib/__init__.py
@@ -15,10 +15,9 @@
 # This module is largely a wrapper around `jaxlib` that performs version
 # checking on import.
 
-import platform
+import gc
 import re
 import os
-import warnings
 from typing import Optional, Tuple
 
 try:
@@ -94,6 +93,11 @@ pytree = xla_client._xla.pytree
 jax_jit = xla_client._xla.jax_jit
 pmap_lib = xla_client._xla.pmap_lib
 
+# XLA garbage collection: see https://github.com/google/jax/issues/14882
+def _xla_gc_callback(*args):
+  xla_client._xla.collect_garbage()
+gc.callbacks.append(_xla_gc_callback)
+
 import jaxlib.gpu_solver as gpu_solver  # pytype: disable=import-error
 import jaxlib.gpu_sparse as gpu_sparse  # pytype: disable=import-error
 import jaxlib.gpu_prng as gpu_prng  # pytype: disable=import-error
diff --git a/tests/api_test.py b/tests/api_test.py
index 978c22547..18e525500 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -9629,5 +9629,18 @@ class AutodidaxTest(jtu.JaxTestCase):
     autodidax_module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(autodidax_module)
 
+class GarbageCollectionTest(jtu.JaxTestCase):
+  def test_xla_gc_callback(self):
+    # https://github.com/google/jax/issues/14882
+    x_np = np.arange(10, dtype='int32')
+    x_jax = jax.device_put(x_np)
+    x_np_weakref = weakref.ref(x_np)
+
+    del x_np
+    del x_jax
+    gc.collect()
+
+    assert x_np_weakref() is None
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax
Issue: `jnp.sort` and `jnp.vdot` raise error under `jax.disable_jit()`

### Description

`jnp.sort`, `jnp.vdot` and a few others in `jax.numpy` will raise under `jax.disable_jit()`. IMO, the fact that exactly *these* functions go wrong is not important, nor is the fact that `jax.disable_jit()` is involved.

Taking a look at their implementation, on can spot a pattern:

```python
@partial(jit, static_argnums=(...))
def f(a: ArrayLike):
  util._check_arraylike('f', a)
  bar = foo(a.ravel())  # << not good
  ...
```

Without jit, the scalar case will get caught by the tests, lead us to the fix naturally. These functions do array conversion implicitly on the jit boundary hides the fact that array methods might not be available.

Fixing them is easy enough, but I want to ask/discuss something more general here:
1. Maybe we should insert these explicit conversions after `_check_arraylike` everywhere, in the form of `asarray` or explicit `reshape` or `ravel`.
2. If the pattern in 1. is desirable, is it always safe to use the simpler and faster variant `util._asarray`? I think if we can, we should.
3. I see similar problems (that JAX maintainers have rather canonical solutions in JAX itself) popping up almost every time someone tries to contribute to `jax.scipy` etc. Can we have a contributor guide documenting these issues? (`array` conversion, `dtype` handling, type hinting, `jnp` vs `lax`, standard ways to deal with static arguments: `concrete_or_error`, etc.)
4. What is the best place to ask these broader questions that are also half bug reports? Should the Github discussions be preferred?

@jakevdp I bore you the most in many related issues before, so I'm gonna tag you here.

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 ed430d46c..7086a4270 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3057,7 +3057,7 @@ def vdot(a, b, *, precision=None):
   util._check_arraylike("vdot", a, b)
   if issubdtype(_dtype(a), complexfloating):
     a = ufuncs.conj(a)
-  return dot(a.ravel(), b.ravel(), precision=precision)
+  return dot(ravel(a), ravel(b), precision=precision)
 
 
 @util._wraps(np.tensordot, lax_description=_PRECISION_DOC)
@@ -3487,9 +3487,9 @@ def sort(a, axis: Optional[int] = -1, kind='quicksort', order=None):
     raise ValueError("'order' argument to sort is not supported.")
 
   if axis is None:
-    return lax.sort(a.ravel(), dimension=0)
+    return lax.sort(ravel(a), dimension=0)
   else:
-    return lax.sort(a, dimension=_canonicalize_axis(axis, ndim(a)))
+    return lax.sort(asarray(a), dimension=_canonicalize_axis(axis, ndim(a)))
 
 @util._wraps(np.sort_complex)
 @jit
Repository: google/jax
Issue: `argmax` and `argmin` doesn't raise when the argument `out` is used.

### Description

https://github.com/google/jax/blob/3fbe4e8230c60e40a1c54732d879438b9d39bd16/jax/_src/numpy/lax_numpy.py#L3401-L3410

`out` is never passed to `_argmax` where the error handling is actually implemented. I guess wrapped functions like `_argmax` shouldn't have default arguments.

Also, `keepdims` should be key word only.

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 ed430d46c..41a79e70b 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3399,14 +3399,15 @@ def argwhere(a, *, size=None, fill_value=None):
 
 
 @util._wraps(np.argmax, skip_params=['out'])
-def argmax(a, axis: Optional[int] = None, out=None, keepdims=None):
-  return _argmax(a, None if axis is None else operator.index(axis), keepdims=bool(keepdims))
-
-@partial(jit, static_argnames=('axis', 'keepdims'), inline=True)
-def _argmax(a, axis: Optional[int] = None, out=None, keepdims=False):
+def argmax(a: ArrayLike, axis: Optional[int] = None, out=None, keepdims=None) -> Array:
   util._check_arraylike("argmax", a)
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.argmax is not supported.")
+  return _argmax(asarray(a), None if axis is None else operator.index(axis),
+                 keepdims=bool(keepdims))
+
+@partial(jit, static_argnames=('axis', 'keepdims'), inline=True)
+def _argmax(a: Array, axis: Optional[int] = None, keepdims: bool = False) -> Array:
   if axis is None:
     dims = list(range(ndim(a)))
     a = ravel(a)
@@ -3419,14 +3420,15 @@ def _argmax(a, axis: Optional[int] = None, out=None, keepdims=False):
   return expand_dims(result, dims) if keepdims else result
 
 @util._wraps(np.argmin, skip_params=['out'])
-def argmin(a, axis: Optional[int] = None, out=None, keepdims=None):
-  return _argmin(a, None if axis is None else operator.index(axis), keepdims=bool(keepdims))
-
-@partial(jit, static_argnames=('axis', 'keepdims'), inline=True)
-def _argmin(a, axis: Optional[int] = None, out=None, keepdims=False):
+def argmin(a: ArrayLike, axis: Optional[int] = None, out=None, keepdims=None) -> Array:
   util._check_arraylike("argmin", a)
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.argmin is not supported.")
+  return _argmin(asarray(a), None if axis is None else operator.index(axis),
+                 keepdims=bool(keepdims))
+
+@partial(jit, static_argnames=('axis', 'keepdims'), inline=True)
+def _argmin(a: Array, axis: Optional[int] = None, keepdims: bool = False) -> Array:
   if axis is None:
     dims = list(range(ndim(a)))
     a = ravel(a)
Repository: google/jax
Issue: Numpy and JAX behaviours differ for sums over scalars.

### Description

When passed a scalar, `jax.numpy` cannot sum over the first nor last axis.

```python
import jax.numpy as jnp
import numpy as np

arr = 0.  # An array that is scalar for some reason
np.sum(arr, 0)  # no pb
np.sum(arr, -1)  # no pb

jnp.sum(arr, 0)  # raises
jnp.sum(arr, -1)  # raises
```

This is also valid for the `jnp.prod` function.
I believe this should likely be a valid input to `jnp.sum` WDYT?

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

jax==0.4.3 jaxlib==0.4.3+cuda11.cudnn86

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

CPU/GPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

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

diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index 712aa4360..37d6d3fcd 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -657,6 +657,7 @@ def _make_cumulative_reduction(np_reduction: Any, reduction: Callable[..., Array
 
     if axis is None or _isscalar(a):
       a = lax.reshape(a, (np.size(a),))
+    if axis is None:
       axis = 0
 
     a_shape = list(np.shape(a))
Repository: google/jax
Issue: ndarray.round() stopped working without an argument.

### Description

Numpy `round()` as a ndarray member function used to work without an argument but it is no longer working.

```python
_ = jnp.zeros(1).round()
```

Expected to execute without an error, but fails with error message

```
TypeError: round() missing required argument 'number' (pos 1)
```

According to Numpy documentation, this should work without an argument.
https://numpy.org/doc/stable/reference/generated/numpy.ndarray.round.html

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

head

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

cpu

### Additional system info

python 3.10, linux

### NVIDIA GPU info

_No response_

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

diff --git a/jax/_src/numpy/array_methods.py b/jax/_src/numpy/array_methods.py
index 364d2f728..576236d0e 100644
--- a/jax/_src/numpy/array_methods.py
+++ b/jax/_src/numpy/array_methods.py
@@ -725,7 +725,7 @@ _array_methods = {
   "ravel": lax_numpy.ravel,
   "repeat": lax_numpy.repeat,
   "reshape": _reshape,
-  "round": round,
+  "round": lax_numpy.round,
   "searchsorted": lax_numpy.searchsorted,
   "sort": lax_numpy.sort,
   "squeeze": lax_numpy.squeeze,
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index b559e78a7..c68162d70 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -858,6 +858,10 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
                         jround(jnp.array(1.234, jnp.float32)),
                         check_dtypes=False)
 
+  def testRoundMethod(self):
+    # https://github.com/google/jax/issues/15190
+    (jnp.arange(3.) / 5.).round()  # doesn't crash
+
   @jtu.sample_product(shape=[(5,), (5, 2)])
   def testOperatorReversed(self, shape):
     rng = jtu.rand_default(self.rng())
Repository: google/jax
Issue: CUDA packages don't have version requirements specified

### Description

Without specified version pip uses whichever versions it has cached. It is fine for now for CUDA 12 since only one cudnn version supports it, but is a problem with CUDA 11 since incompatible version may be present in package cache.

### What jax/jaxlib version are you using?

0.4.7

Changed files:
setup.py

diff --git a/setup.py b/setup.py
index 2214d11d6..2ef93a56c 100644
--- a/setup.py
+++ b/setup.py
@@ -96,13 +96,13 @@ setup(
 
         'cuda11_pip': [
           f"jaxlib=={_current_jaxlib_version}+cuda11.cudnn{_default_cuda11_cudnn_version}",
-          "nvidia-cublas-cu11",
-          "nvidia-cuda-nvcc-cu11",
-          "nvidia-cuda-runtime-cu11",
-          "nvidia-cudnn-cu11",
-          "nvidia-cufft-cu11",
-          "nvidia-cusolver-cu11",
-          "nvidia-cusparse-cu11",
+          "nvidia-cublas-cu11>=11.11",
+          "nvidia-cuda-nvcc-cu11>=11.8",
+          "nvidia-cuda-runtime-cu11>=11.8",
+          "nvidia-cudnn-cu11>=8.6",
+          "nvidia-cufft-cu11>=10.9",
+          "nvidia-cusolver-cu11>=11.4",
+          "nvidia-cusparse-cu11>=11.7",
         ],
 
         'cuda12_pip': [
Repository: google/jax
Issue: Add missing exports to jax.config and jax.debug

### Description

Add missing exports to [jax.config](https://github.com/google/jax/blob/main/jax/config.py) and [jax.debug](https://github.com/google/jax/blob/main/jax/debug.py) (and any other such files). See #7570 for context.

Example:
```python3
from jax.config import config
from jax.debug import print, callback, breakpoint
```

mypy yields
```
example.py:1: error: Module "jax.config" does not explicitly export attribute "config"  [attr-defined]
example.py:2: error: Module "jax.debug" does not explicitly export attribute "print"  [attr-defined]
example.py:2: error: Module "jax.debug" does not explicitly export attribute "callback"  [attr-defined]
example.py:2: error: Module "jax.debug" does not explicitly export attribute "breakpoint"  [attr-defined]
```

pyright yields
```
  /Users/carlos/Desktop/example.py:1:24 - error: "config" is not exported from module "jax.config"
    Import from "jax._src.config" instead (reportPrivateImportUsage)
  /Users/carlos/Desktop/example.py:2:23 - error: "print" is not exported from module "jax.debug"
    Import from "jax._src.debugging" instead (reportPrivateImportUsage)
  /Users/carlos/Desktop/example.py:2:30 - error: "callback" is not exported from module "jax.debug"
    Import from "jax._src.debugging" instead (reportPrivateImportUsage)
  /Users/carlos/Desktop/example.py:2:40 - error: "breakpoint" is not exported from module "jax.debug"
    Import from "jax._src.debugger.core" instead (reportPrivateImportUsage)
```

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

jax 0.4.8, jaxlib 0.4.7

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

CPU

### Additional system info

Python 3.11.2, macOS 11.7.4

### NVIDIA GPU info

_No response_

Changed files:
jax/debug.py

diff --git a/jax/debug.py b/jax/debug.py
index 276f669ff..64578524d 100644
--- a/jax/debug.py
+++ b/jax/debug.py
@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+__all__ = ["callback", "print", "DebugEffect", "visualize_array_sharding",
+           "inspect_array_sharding", "visualize_sharding", "breakpoint"]
 from jax._src.debugging import debug_callback as callback
 from jax._src.debugging import debug_print as print
 from jax._src.debugging import DebugEffect
Repository: google/jax
Issue: bad error on VJP of functions returning typed key arrays

Found by @NeilGirdhar in #14046!

```python
from jax import enable_custom_prng, vjp
from jax.random import PRNGKey

with enable_custom_prng():
    def f(i):
        return PRNGKey(i)

    out, f_vjp = vjp(f, 1)  # Fails!
```
Error is:
```
TypeError: Value 'PRNGKeyArray[fry] { [0 1] }' with dtype key<fry> is not a valid JAX array type. Only arrays of numeric types are supported by JAX.
```

cc #9263

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

diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
index 42d73877e..af9512dcd 100644
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -617,10 +617,10 @@ def dtype(x: Any, *, canonicalize: bool = False) -> DType:
       dt = np.result_type(x)
     except TypeError as err:
       raise TypeError(f"Cannot determine dtype of {x}") from err
-  if dt not in _jax_dtype_set:
+  if dt not in _jax_dtype_set and not core.is_opaque_dtype(dt):
     raise TypeError(f"Value '{x}' with dtype {dt} is not a valid JAX array "
                     "type. Only arrays of numeric types are supported by JAX.")
-  return canonicalize_dtype(dt) if canonicalize else dt
+  return canonicalize_dtype(dt, allow_opaque_dtype=True) if canonicalize else dt
 
 def _lattice_result_type(*args: Any) -> Tuple[DType, bool]:
   dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
diff --git a/tests/random_test.py b/tests/random_test.py
index 6a629c0ee..bb056116d 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -526,6 +526,11 @@ class PrngTest(jtu.JaxTestCase):
     key = random.PRNGKey(0)
     self.assertIsInstance(key, jax.Array)
 
+  def test_key_output_vjp(self):
+    # See https://github.com/google/jax/issues/14856
+    def f(seed): return random.PRNGKey(seed)
+    jax.vjp(f, 1)  # doesn't crash
+
 
 class ThreefryPrngTest(jtu.JaxTestCase):
   def test_seed_no_implicit_transfers(self):
Repository: google/jax
Issue: forward mode differentiation issue with scipy.special.xlogy

### Description

A few days ago, a [commit](https://github.com/google/jax/commit/dd023e266e6616494cdd590959103ddc646109c4) to jax.scipy.special seems to be causing my previously-working code to bomb out with NaN errors.  I believe the root problem is that you've added a custom jvp to xlogy and xlog1py which need the same safety features as the original functions.

```@custom_derivatives.custom_jvp
@_wraps(osp_special.xlogy, module='scipy.special')
def xlogy(x: ArrayLike, y: ArrayLike) -> Array:
  # Note: xlogy(0, 0) should return 0 according to the function documentation.
  x, y = promote_args_inexact("xlogy", x, y)
  x_ok = x != 0.
  safe_x = jnp.where(x_ok, x, 1.)
  safe_y = jnp.where(x_ok, y, 1.)
  return jnp.where(x_ok, lax.mul(safe_x, lax.log(safe_y)), jnp.zeros_like(x))

def _xlogy_jvp(primals, tangents):
  (x, y) = primals
  (x_dot, y_dot) = tangents
  result = xlogy(x, y)
  return result, (x_dot * lax.log(y) + y_dot * x / y).astype(result.dtype)
xlogy.defjvp(_xlogy_jvp)
```
My code is obtaining NaNs in _xlogy_jvp when I call scipy.special.entr(p) which then calls xlogy(p,p) and one of the elements of p is zero.

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

jax v0.4.9

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

TPU

### Additional system info

Colab

### NVIDIA GPU info

_No response_

Changed files:
tests/lax_scipy_test.py, jax/_src/scipy/special.py

diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
index fb33d9e92..d8bedc771 100644
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -155,12 +155,24 @@ def _xlog1py_jvp(primals, tangents):
   return result, (x_dot * lax.log1p(y) + y_dot * x / (1 + y)).astype(result.dtype)
 xlog1py.defjvp(_xlog1py_jvp)
 
+@custom_derivatives.custom_jvp
+def _xlogx(x):
+  """Compute x log(x) with well-defined derivatives."""
+  return xlogy(x, x)
+
+def _xlogx_jvp(primals, tangents):
+  x, = primals
+  x_dot, = tangents
+  return  _xlogx(x), x_dot * (lax.log(x) + 1)
+_xlogx.defjvp(_xlogx_jvp)
+
+
 @_wraps(osp_special.entr, module='scipy.special')
 def entr(x: ArrayLike) -> Array:
   x, = promote_args_inexact("entr", x)
   return lax.select(lax.lt(x, _lax_const(x, 0)),
                     lax.full_like(x, -np.inf),
-                    lax.neg(xlogy(x, x)))
+                    lax.neg(_xlogx(x)))
 
 
 @_wraps(osp_special.multigammaln, update_doc=False)
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
index 3ecf24e67..3e8c84398 100644
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -243,6 +243,11 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
 
     jtu.check_grads(lsp_special.xlog1py, (x0, y0), order=2)
 
+  def testGradOfEntrAtZero(self):
+    # https://github.com/google/jax/issues/15709
+    self.assertEqual(jax.jacfwd(lsp_special.entr)(0.0), jnp.inf)
+    self.assertEqual(jax.jacrev(lsp_special.entr)(0.0), jnp.inf)
+
   @jtu.sample_product(
     [dict(order=order, z=z, n_iter=n_iter)
      for order, z, n_iter in zip(