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

Issue:
autodidax fails with jaxlib 0.4.2
After updating to jaxlib 0.4.2, `python docs/autodidax.py` fails on this line: https://github.com/google/jax/blob/d2845f264f283897ad9cb8cf742fcd41bd303f5b/docs/autodidax.py#L1570

with the following traceback:
```pytb
Traceback (most recent call last):
  File "docs/autodidax.py", line 1675, in <module>
    z = f(3., 4.)  # 'tracing!' prints the first time
  File "docs/autodidax.py", line 1510, in f_jitted
    outs = bind(xla_call_p, *consts, *args, jaxpr=jaxpr, num_consts=len(consts))
  File "docs/autodidax.py", line 327, in bind
    outs = top_trace.process_primitive(prim, tracers, params)
  File "docs/autodidax.py", line 412, in process_primitive
    return impl_rules[primitive](*tracers, **params)
  File "docs/autodidax.py", line 1554, in xla_call_impl
    execute = xla_callable(IDHashable(jaxpr), hashable_consts)
  File "docs/autodidax.py", line 1570, in xla_callable
    compiled = xb.get_backend(None).compile(c.build(out))
TypeError: compile(): incompatible function arguments. The following argument types are supported:
    1. (self: jaxlib.xla_extension.Client, computation: str, compile_options: jaxlib.xla_extension.CompileOptions = <jaxlib.xla_extension.CompileOptions object at 0x7f84b1946330>, host_callbacks: List[capsule] = []) -> StatusOr[xla::PyLoadedExecutable]

Invoked with: <jaxlib.xla_extension.Client object at 0x7f84b3a6cbf0>, <jaxlib.xla_extension.XlaComputation object at 0x7f84b3a727f0>
```


Discussion:

Resolving patch:
diff --git a/docs/autodidax.py b/docs/autodidax.py
--- a/docs/autodidax.py
+++ b/docs/autodidax.py
@@ -1567,7 +1567,8 @@ def xla_callable(hashable_jaxpr: IDHashable,
   xla_params = _xla_params(c, in_avals)
   outs = jaxpr_subcomp(c, jaxpr, xla_consts + xla_params)
   out = xops.Tuple(c, outs)
-  compiled = xb.get_backend(None).compile(c.build(out))
+  compiled = xb.get_backend(None).compile(
+    xc._xla.mlir.xla_computation_to_mlir_module(c.build(out)))
   return partial(execute_compiled, compiled, [v.aval for v in jaxpr.outs])
 
 def _xla_consts(c: xe.XlaBuilder, consts: List[Any]) -> List[xe.XlaOp]:


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -20,7 +20,9 @@
 import enum
 from functools import partial
 import inspect
+import importlib
 import operator
+import os
 import platform
 import re
 import subprocess
@@ -9509,5 +9511,17 @@ def test_print_environment_info(self, return_string):
     assert f"jaxlib: {lib.version_str}" in result
     assert f"numpy:  {np.__version__}" in result
 
+class AutodidaxTest(jtu.JaxTestCase):
+  def test_autodidax_smoketest(self):
+    autodidax_file = os.path.join(
+      os.path.dirname(os.path.dirname(__file__)),
+      'docs',
+      'autodidax.py')
+    if not os.path.exists(autodidax_file):
+      self.skipTest("Cannot locate autodidax.py")
+    spec = importlib.util.spec_from_file_location('autodidax', 'docs/autodidax.py')
+    autodidax_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(autodidax_module)
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax

Issue:
Implement `jnp.tril_indices`, `jnp.triu_indices` etc. in JAX.
Implement `jnp.tril_indices` etc. using XLA primitives instead of wrapping `numpy` implementations, so we don't materialise their output while staging.


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2673,11 +2673,35 @@ def wrapper(*args, **kwargs):
     return tuple(asarray(x) for x in f(*args, **kwargs))
   return wrapper
 
-tril_indices = _wrap_indices_function(np.tril_indices)
-triu_indices = _wrap_indices_function(np.triu_indices)
 mask_indices = _wrap_indices_function(np.mask_indices)
 
 
+def _triu_size(n, m, k):
+  if k < 0:
+    return n * m - _triu_size(m, n, (1 - k))
+  elif k >= m:
+    return 0
+  else:
+    mk = _min(n, m - k)
+    return mk * (mk + 1) // 2 + mk * (m - k - mk)
+
+
+@_wraps(np.triu_indices)
+def triu_indices(n, k, m=None):
+  n = core.concrete_or_error(operator.index, n, "n argument of jnp.triu_indices")
+  k = core.concrete_or_error(operator.index, k, "k argument of jnp.triu_indices")
+  m = n if m is None else core.concrete_or_error(operator.index, m, "m argument of jnp.triu_indices")
+  return nonzero(triu(ones((n, m)), k=k), size=_triu_size(n, m, k))
+
+
+@_wraps(np.tril_indices)
+def tril_indices(n, k, m=None):
+  n = core.concrete_or_error(operator.index, n, "n argument of jnp.triu_indices")
+  k = core.concrete_or_error(operator.index, k, "k argument of jnp.triu_indices")
+  m = n if m is None else core.concrete_or_error(operator.index, m, "m argument of jnp.triu_indices")
+  return nonzero(tril(ones((n, m)), k=k), size=_triu_size(m, n, -k))
+
+
 @_wraps(np.triu_indices_from)
 def triu_indices_from(arr: ArrayLike, k: int = 0) -> Tuple[Array]:
   arr_shape = shape(arr)


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1992,9 +1992,9 @@ def testTriLU(self, dtype, shape, op, k):
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @jtu.sample_product(
-    n=range(1, 5),
-    k=[-1, 0, 1],
-    m=range(1, 5),
+    n=range(5),
+    k=range(-3, 3),
+    m=[None, *range(5)],
   )
   def testTrilIndices(self, n, k, m):
     np_fun = lambda n, k, m: np.tril_indices(n, k=k, m=m)
@@ -2003,9 +2003,9 @@ def testTrilIndices(self, n, k, m):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
 
   @jtu.sample_product(
-    n=range(1, 5),
-    k=[-1, 0, 1],
-    m=range(1, 5),
+    n=range(5),
+    k=range(-3, 3),
+    m=[None, *range(5)],
   )
   def testTriuIndices(self, n, k, m):
     np_fun = lambda n, k, m: np.triu_indices(n, k=k, m=m)
Repository: google/jax

Issue:
Improved documentation on `jax.random.gamma`
Please:

- [x] Check for duplicate requests.
- [x] Describe your goal, and if possible provide a code snippet with a motivating example.

I think the documentation for the `jax.random.gamma` random variable (sampler) could possibly be clearer. In particular, letting the user know that the standard form is implemented and how to incorporate the rate (or alternatively scale) parameter in their samples. Maybe an addendum along the lines of 

*"we implement the standard gamma density. Dividing the sample by the rate is equivalent to sampling from* $\text{Gamma}(\alpha, \text{rate})$*. And multiplying the sample by the scale is equivalent to sampling from* $\text{Gamma}(\alpha, \text{scale})$."

The discussion for this started [here](https://github.com/google/jax/pull/551) and I'm mentioning @mattjj for visibility. Thanks!


Discussion:

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1068,6 +1068,11 @@ def gamma(key: KeyArray,
           dtype: DTypeLikeFloat = dtypes.float_) -> Array:
   """Sample Gamma random values with given shape and float dtype.
 
+  This implements the standard gamma density, with a unit scale/rate parameter.
+  Dividing the sample output by the rate is equivalent to sampling from
+  *gamma(a, rate)*, and multiplying the sample output by the scale is equivalent
+  to sampling from *gamma(a, scale)*.
+
   Args:
     key: a PRNG key used as the random key.
     a: a float or array of floats broadcast-compatible with ``shape``


Test patch:
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
```


Discussion:
From @hawkinsp, it looks like triggering XLA's internal garbage cleanup is the solution here:
```python
import sys
import jax
import numpy as np
from jax._src.lib import xla_client

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

xla_client._xla.collect_garbage()
print(sys.getrefcount(x))
# 2
```
This internal garbage collection is called frequently within XLA when dispatching JAX operations (for example, replacing `_xla.collect_garbage()` with a JAX call like `jax.jit(lambda a: a)(1.0)` shows the same result), but @mattjj suggested we might also add it to [`gc.callbacks`](https://docs.python.org/3/library/gc.html#gc.callbacks) so that when the user forces Python garbage collection, it triggers XLA garbage collection as well.
With this addition, `gc.collect()` causes the reference counts to be as expected.
```python
import gc
from jax._src.lib import xla_client
gc.callbacks.append(lambda *args: xla_client._xla.collect_garbage())
```

Resolving patch:
diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
--- 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 @@ def _parse_version(v: str) -> Tuple[int, ...]:
 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


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -9629,5 +9629,18 @@ def test_autodidax_smoketest(self):
     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:
`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.


Discussion:
Thanks for the report!

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -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)


Test patch:
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_


Discussion:
Thanks for the report – I don't see this behavior documented anywhere in numpy, so I suspect it's an implementation detail rather than intended design. I think JAX's behavior makes more sense: in general if you have an array `a` and you sum over `axis=a.ndim`, you will get an error. I think this also makes sense when `ndim = 0`. What do you think?
@jakevdp Should we also raise for `cumsum` and friends

```python
jnp.cumsum(0, axis=2)
# Array([0], dtype=int32, weak_type=True)
```
The NumPy version behaviours slightly differently, and raises for `axis not in {0, -1}`. Both versions violates the documentation

> The result has the same size as `a`, and the same shape as `a` if `axis` is not None or `a` is a 1-d array.

Sometimes NumPy's semantics is just 🤦 .
That looks like a bug in the implementation of `jnp.cumsum`. Are you interested in putting together a fix?
> Are you interested in putting together a fix?

@jakevdp Sure.
I've checked other stats function. For instance `np.mean` raises on scalars when an axis is passed. It seems to confirm that this is more of an implementation side effect than a feature of NumPy.

Resolving patch:
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -657,6 +657,7 @@ def _cumulative_reduction(a: ArrayLike, axis: Axis = None,
 
     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))


Test patch:
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_


Discussion:
Probably related to https://github.com/google/jax/pull/15172

Resolving patch:
diff --git a/jax/_src/numpy/array_methods.py b/jax/_src/numpy/array_methods.py
--- a/jax/_src/numpy/array_methods.py
+++ b/jax/_src/numpy/array_methods.py
@@ -725,7 +725,7 @@ def max(self, values, *, indices_are_sorted=False, unique_indices=False,
   "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,


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -858,6 +858,10 @@ def testOperatorRound(self, jit):
                         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:
Rayleigh sampling doesn't support broadcasting
### Description

```python
import jax.numpy as jnp
from jax.random import rayleigh, PRNGKey

key = PRNGKey(123)
sigma = jnp.ones(10)
rayleigh(key, sigma, (3, 10))
```

I'm happy to submit a PR for this.


Discussion:
Thanks - the entire random module treats shape and dtype arguments inconsistently, and needs to be overhauled. I have some work in progress to do this more comprehensively (first PR is #15177 and addresses dtype handling).

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -2036,6 +2036,7 @@ def _rayleigh(key, scale, shape, dtype) -> Array:
     _check_shape("rayleigh", shape, np.shape(scale))
   u = uniform(key, shape, dtype)
   scale = scale.astype(dtype)
+  scale = jnp.broadcast_to(scale, shape)
   log_u = lax.log(u)
   n_two = _lax_const(scale, -2)
   sqrt_u = lax.sqrt(lax.mul(log_u, n_two))


Test patch:
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


Discussion:

Resolving patch:
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
--- 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))


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -526,6 +526,11 @@ def test_isinstance(self):
     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:
`jax.numpy.asarray` does not function under autodiff
### Description

```python
import jax
p, t = jax.jvp(jax.numpy.asarray, (1.,), (2.,))
print(p, type(p))
print(t, type(t))
# 1.0 <class 'float'>
# 2.0 <class 'float'>
```
In particular the primal computation produces a different type than is produced without autodiff, via just `jax.numpy.asarray(1.)`.

### What jax/jaxlib version are you using?

JAX 0.4.8


Discussion:
Thanks for noticing this and reporting it so clearly!
For anyone else still getting bit by this, my workaround is pretty simple:
```python
@jax.custom_jvp
def fixed_asarray(x):
    return jnp.asarray(x)

@fixed_asarray.defjvp
def _fixed_asarray_jvp(x, tx):
    (x,) = x
    (tx,) = tx
    return fixed_asarray(x), fixed_asarray(tx)
```
I just hit this in refactoring `jnp.dot` in #16721. @mattjj - any plans on this?

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -542,7 +542,10 @@ def _convert_element_type(operand: ArrayLike, new_dtype: Optional[DTypeLike] = N
     operand = np.asarray(operand).astype(new_dtype)
     old_weak_type = False
 
-  if (old_dtype, old_weak_type) == (new_dtype, weak_type) and isinstance(operand, Array):
+  if ((old_dtype, old_weak_type) == (new_dtype, weak_type) and
+      isinstance(operand, Array) and
+      not (isinstance(operand, core.Tracer) and
+           isinstance(core.get_aval(operand), core.ConcreteArray))):
     return type_cast(Array, operand)
   else:
     return convert_element_type_p.bind(operand, new_dtype=new_dtype,


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4358,6 +4358,12 @@ def test_jit_custom_floats(self, dtype):
     args_maker = lambda: [jnp.ones((), dtype=dtype)]
     self._CompileAndCheck(f, args_maker)
 
+  def test_jvp_asarray_returns_array(self):
+    # https://github.com/google/jax/issues/15676
+    p, t = jax.jvp(jax.numpy.asarray, (1.,), (2.,))
+    _check_instance(self, p)
+    _check_instance(self, t)
+
 
 class RematTest(jtu.JaxTestCase):