anon-iclr-submission/benchname-bug-localization huggingface dataset:
Repository: google/jax
Issue: `jit` + `linear_transpose` raises error due to missing DCE  call

### Description

```python
import jax

def f1(x):
    return 2 * x, x > 0

def g1(x):
    return f1(x)[0]

jax.linear_transpose(g1, 1.)(1.)  # works

@jax.jit
def f2(x):
    return 2 * x, x > 0

def g2(x):
    return f2(x)[0]

jax.linear_transpose(g2, 1.)(1.)  # error!
# jax._src.traceback_util.UnfilteredStackTrace: NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'gt' not implemented
```

It looks like `jax.linear_transpose` is missing a call to `pe.dce_jaxpr` prior to attempting the transpose.

I checked that `jax.grad` doesn't have the same issue. I assume it has a DCE call in there already somewhere.

### What jax/jaxlib version are you using?

JAX 0.4.8

Changed files:
tests/api_test.py, jax/_src/api.py

diff --git a/jax/_src/api.py b/jax/_src/api.py
index 15e9d07d5..43341b8fc 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2250,6 +2250,7 @@ def linear_transpose(fun: Callable, *primals, reduce_axes=()) -> Callable:
   in_pvals = map(pe.PartialVal.unknown, in_avals)
   jaxpr, out_pvals, const = pe.trace_to_jaxpr_nounits(flat_fun, in_pvals,
                                                       instantiate=True)
+  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)
   out_avals, _ = unzip2(out_pvals)
   out_dtypes = map(dtypes.dtype, out_avals)
   if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)
diff --git a/tests/api_test.py b/tests/api_test.py
index f788bb097..4bbde83c5 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2076,6 +2076,12 @@ class APITest(jtu.JaxTestCase):
     expected = 6
     self.assertEqual(actual, expected)
 
+  def test_linear_transpose_dce(self):
+    # https://github.com/google/jax/issues/15660
+    f = jit(lambda x: (2 * x, x > 0))
+    g = lambda x: f(x)[0]
+    api.linear_transpose(g, 1.)(1.)
+
   def test_linear_transpose_error(self):
     with self.assertRaisesRegex(
         TypeError, "linear_transpose only supports"):
Repository: google/jax
Issue: Inconsistent dtype output in jax.random.choice for double precision enabled

### Description

When using double precision, `jax.random.choice` returns inconsistent dtypes. When using a weight vector, the output dtype is set using `searchsorted` rules, i.e., based on the size of the array of weights at hand, while when none is used (uniform weights), the dtype is set using `jax.random.randint` rules (defaulting to `jnp.int_`). I believe the behaviour should be homogenous, no matter what the choice ends up being.

```python
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)
key = jax.random.PRNGKey(42)
w = jnp.array([1/3, 1/3, 1/3])

print(jax.random.choice(key, 3).dtype)
print(jax.random.choice(key, 3, p=w).dtype)
```

In practice, I believe that the implementation choice made in https://github.com/google/jax/blob/c22302ae2e0913efd5be23e136f2ff14b520d9cf/jax/_src/numpy/lax_numpy.py#L4824 is perhaps a bit redundant. Indeed, if the config has not enables double precision, `int64` will default to `int32` anyway, while, if it has, you were happy to pay the price in the first place.

WDYT?



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

jax v0.4.10 

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

CPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

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

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 08d64dadc..56cc53cdf 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -567,7 +567,7 @@ def choice(key: KeyArray,
     if replace:
       p_cuml = jnp.cumsum(p_arr)
       r = p_cuml[-1] * (1 - uniform(key, shape, dtype=p_cuml.dtype))
-      ind = jnp.searchsorted(p_cuml, r)
+      ind = jnp.searchsorted(p_cuml, r).astype(int)
     else:
       # Gumbel top-k trick: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/
       g = -gumbel(key, (n_inputs,), dtype=p_arr.dtype) - jnp.log(p_arr)
diff --git a/tests/random_test.py b/tests/random_test.py
index 26a544688..93fc96c01 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -754,8 +754,10 @@ class LaxRandomTest(jtu.JaxTestCase):
     sample = rand(key, x)
     if not is_range:
       self.assertEqual(dtype, sample.dtype)
-    np_shape = np.shape(np_choice(x, shape or None, replace, p, axis))
-    self.assertEqual(np_shape, sample.shape)
+    expected_shape = np.shape(np_choice(x, shape or None, replace, p, axis))
+    self.assertEqual(expected_shape, sample.shape)
+    expected_dtype = dtypes.result_type(int if is_range else x)
+    self.assertEqual(expected_dtype, sample.dtype)
     if not replace and shape:
       def lsort(x):
         if not math.prod(x.shape): return x
Repository: google/jax
Issue: Potential typo for `jax.log_compiles`

### Description

`jax.log_compiles()` currently sets the flag for `jax_log_checkpoint_residuals` instead of `jax_log_compiles`, which looks like a typo.

This looks like a copy-paste error introduced by 6b4262d9f6ee36aea16ba7cfac639904b499661c.

https://github.com/google/jax/blob/9da52e890537c507b4d559ee6e3d8e8358e77f3f/jax/_src/config.py#L737-L750

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

jax v0.4.10, jaxlib v0.4.10

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

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

Changed files:
jax/_src/config.py

diff --git a/jax/_src/config.py b/jax/_src/config.py
index 68d448111..b09e3084d 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -742,7 +742,7 @@ log_compiles = config.define_bool_state(
           'option is set, the log level is WARNING; otherwise the level is '
           'DEBUG.'))
 
-log_compiles = config.define_bool_state(
+log_checkpoint_residuals = config.define_bool_state(
     name='jax_log_checkpoint_residuals',
     default=False,
     help=('Log a message every time jax.checkpoint (aka jax.remat) is '
Repository: google/jax
Issue: [shard-map] Check failed using unsafe_rbg_key

Based on code from @jheek :

```python
import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'

from functools import partial

import numpy as np

import jax
import jax.numpy as jnp
from jax import random
from jax.experimental.shard_map import shard_map, PartitionSpec as P, Mesh

from flax import linen

jax.config.update('jax_enable_custom_prng', True)  # doesn't crash if we exclude this line!

dense = linen.Dense(128)

mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('model', 'data'))

@jax.jit
def f(rng):
  @partial(shard_map, mesh=mesh, in_specs=P("model"), out_specs=P("model"), check_rep=False)
  def g(rng):
    rng = rng[0]
    return dense.init(rng, jnp.zeros((8, 128)))

  return g(random.split(rng, 4))

f(random.PRNGKey(0)) # ==> crashes VM on ASSERT(IsManual(sharding))
```

Changed files:
jax/_src/interpreters/mlir.py, jax/_src/prng.py, tests/shard_map_test.py

diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 8dde30217..d6d0196ef 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1500,7 +1500,7 @@ def _wrap_with_spmd_op(name: str,
   else:
     backend_config = ""
   result_type = aval_to_ir_type(aval_out)
-  out_shape = aval_out.shape  # type: ignore
+  out_shape = core.physical_aval(aval_out).shape  # type: ignore
   if core.is_constant_shape(out_shape):
     result_shapes = None
   else:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
index d23283eec..a62307c80 100644
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -470,7 +470,8 @@ class KeyTyRules:
       op_sharding_proto = op_sharding_proto.to_proto()
     new_op_sharding = op_sharding_proto.clone()
     tad = list(new_op_sharding.tile_assignment_dimensions)
-    tad.extend([1] * len(key_shape))
+    suffix = [tad.pop()] if op_sharding_proto.replicate_on_last_tile_dim else []
+    tad.extend([1] * len(key_shape) + suffix)
     new_op_sharding.tile_assignment_dimensions = tad
     return xc.HloSharding.from_proto(new_op_sharding)
 
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
index 6b8aaf1e3..5f28ce454 100644
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -740,6 +740,20 @@ class ShardMapTest(jtu.JaxTestCase):
     # error!
     jax.jit(g)(x)  # doesn't crash
 
+  def test_key_array_with_replicated_last_tile_dim(self):
+    # See https://github.com/google/jax/issues/16137
+
+    mesh = jtu.create_global_mesh((2, 4), ('i', 'j'))
+
+    def f(rng):
+      @partial(shard_map, mesh=mesh, in_specs=P('i'), out_specs=P('i'),
+               check_rep=False)
+      def g(rng):
+        return jnp.array([jax.random.normal(rng[0])])
+      return g(jax.random.split(rng, 4))
+
+    jax.jit(f)(jax.random.key(0))  # doesn't crash
+
   # same method appears in api_test.py:DCETest
   # TODO(mattjj): consider moving this method to be a helper in jtu
   def assert_dce_result(self, jaxpr: core.Jaxpr, used_outputs: List[bool],
Repository: google/jax
Issue: Zero-length indices in `jax.numpy.mgrid`

### Description

```python
np.mgrid[()]  # yields an integer type array of zero length
jnp.mgrid[()]  # raises a ValueError because it tries to stack a tuple of zero length
```

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

jax commit f729da4a368e9, jaxlib 0.4.1

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

CPU/GPU

### Additional system info

python 3.10.9, Linux (Arch Linux)

### NVIDIA GPU info

```
| NVIDIA-SMI 525.60.13    Driver Version: 525.60.13    CUDA Version: 12.0     |              
|-------------------------------+----------------------+----------------------+              
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |              
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |              
|                               |                      |               MIG M. |              
|===============================+======================+======================|              
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0 Off |                  N/A |              
|  0%   40C    P8    25W / 350W |     17MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
```

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

diff --git a/jax/_src/numpy/index_tricks.py b/jax/_src/numpy/index_tricks.py
index 5fff1df13..a61fa2749 100644
--- a/jax/_src/numpy/index_tricks.py
+++ b/jax/_src/numpy/index_tricks.py
@@ -54,7 +54,11 @@ class _IndexGrid(abc.ABC):
     with jax.numpy_dtype_promotion('standard'):
       output = _promote_dtypes(*output)
     output_arr = jnp.meshgrid(*output, indexing='ij', sparse=self.sparse)
-    return output_arr if self.sparse else jnp.stack(output_arr, 0)
+    if self.sparse:
+      return output_arr
+    if len(output_arr) == 0:
+      return jnp.arange(0)
+    return jnp.stack(output_arr, 0)
 
 
 class _Mgrid(_IndexGrid):
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 678562b71..72475e96f 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4353,6 +4353,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     # wrap indexer for appropriate dtype defaults.
     np_mgrid = _indexer_with_default_outputs(np.mgrid)
     assertAllEqual = partial(self.assertAllClose, atol=0, rtol=0)
+    assertAllEqual(np_mgrid[()], jnp.mgrid[()])
     assertAllEqual(np_mgrid[:4], jnp.mgrid[:4])
     assertAllEqual(np_mgrid[:4,], jnp.mgrid[:4,])
     assertAllEqual(np_mgrid[:4], jax.jit(lambda: jnp.mgrid[:4])())
Repository: google/jax
Issue: `stax.BatchNorm` incorrect normalization

jax.experimental.stax on current master, line 140 ( `BatchNorm` function):
```python
z = (x - mean) / (var + epsilon)**2
```
The square in the denominator should probably be a square root.

Changed files:
jax/experimental/stax.py

diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
index 311c9d3f5..caff9a984 100644
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -137,7 +137,7 @@ def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
   def apply_fun(params, x, rng=None):
     beta, gamma = params
     mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True)
-    z = (x - mean) / (var + epsilon)**2
+    z = (x - mean) / np.sqrt(var + epsilon)
     if center and scale: return gamma * z + beta
     if center: return z + beta
     if scale: return gamma * z
Repository: google/jax
Issue: Error differentiating QR decomposition

This appears to be a bug but unfortunately I don't have time to delve into it.

Below is a simple script to do gradient descent on a function involving a QR decomposition. It fails when I execute `training_gradient_fun` with `TypeError: No abstraction handler for type: <class 'jax.ad_util.Zero'>`.

Ubuntu 16.04
python 3.6.6
jax 0.1.20
jaxlib 0.1.9

```
import jax
import jax.numpy as np
import numpy as onp

def loss(weights, inputs):
    q, r = np.linalg.qr(weights)
    a = np.matmul(q, weights)
    return np.mean(a ** 2)

training_gradient_fun = jax.jit(jax.grad(loss))

inputs = onp.random.rand(5, 5)
weights = onp.random.rand(5, 5)

for i in range(100):
    weights -= 0.1 * training_gradient_fun(weights, inputs)
```

Changed files:
jax/lax.py

diff --git a/jax/lax.py b/jax/lax.py
index 3ffc2598f..556886e99 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1363,7 +1363,7 @@ ad.primitive_transposes[add_p] = _add_transpose
 
 def _sub_transpose(t, x, y):
   assert x is None and y is None  # computation must be linear, not affine
-  return [t, neg(t)]
+  return [t, neg(t) if t is not ad_util.zero else ad_util.zero]
 
 sub_p = standard_binop([_num, _num], 'sub')
 ad.defjvp(sub_p,
Repository: google/jax
Issue: bug in np.full / lax.full for DeviceArray wrapped scalars

If you pass in a jax-wrapped scalar as the fill-constant value to np.fill, it breaks, e.g.:
```python
np.full((1,1), np.array([1])[0])
```
this hits the condition in lax.full:
```python
 elif isinstance(fill_value, xla.DeviceValue):
    return FilledConstant(convert_element_type(fill_value, dtype), shape)
```
no matter what dtype is, the convert_element_type result is still a DeviceArray, then
this line in lax.FillConstant:
```python
class FilledConstant(xla.DeviceConstant):
  __slots__ = ["fill_value"]

  def __init__(self, fill_value, shape):
    assert type(fill_value) is onp.ndarray
```
wants the fill value to be an onp ndarray instead.

I'm not too certain about the type logic here to make a PR confidently.

Changed files:
jax/lax.py, tests/lax_numpy_test.py

diff --git a/jax/lax.py b/jax/lax.py
index b7b8d60d4..444c86b44 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -462,7 +462,8 @@ def full(shape, fill_value, dtype):
   if onp.isscalar(fill_value) or type(fill_value) is onp.ndarray:
     return FilledConstant(onp.asarray(fill_value, dtype), shape)
   elif isinstance(fill_value, xla.DeviceValue):
-    return FilledConstant(convert_element_type(fill_value, dtype), shape)
+    val = onp.asarray(fill_value, dtype)
+    return FilledConstant(val, shape)
   else:
     return broadcast(convert_element_type(fill_value, dtype), shape)
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index d98c3dd58..8902380a2 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1187,7 +1187,6 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       self._CheckAgainstNumpy(lnp_op, onp_op, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True)
 
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_n={}_increasing={}".format(
           jtu.format_shape_dtype_string([shape], dtype),
@@ -1207,6 +1206,10 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=False)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=False)
 
+  def testIssue330(self):
+    x = lnp.full((1, 1), lnp.array([1])[0])  # doesn't crash
+    self.assertEqual(x[0, 0], 1)
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax
Issue: NumPy 1.16.0 causes warnings

```jax/lax.py:769: FutureWarning: Conversion of the second argument of issubdtype from complex to np.complexfloating is deprecated. In future, it will be treated as np.complex128 == np.dtype(complex).type.```

(I don't have code to reproduce it right now, because any attempt to create a minimal example makes the warning disappear even though it doesn't make any sense.)

Changed files:
jax/lax.py

diff --git a/jax/lax.py b/jax/lax.py
index 8aec41764..471067889 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -831,7 +831,7 @@ def _brcast_to(x, shape):
 
 _f32 = {onp.float32}
 _float = {onp.floating}
-_complex = {onp.complex}
+_complex = {onp.complexfloating}
 _complex_elem_types = {onp.float32, onp.float64}
 _int = {onp.integer}
 _bool = {onp.bool_}
Repository: google/jax
Issue: jax.random.randint range must be valid

`random.randint` doesn't check that `minval` < `maxval`. If it's not, you get surprising results. For example,
`print(random.randint(random.PRNGKey(10003), (), 0, 0))`
yields -1234647498.

The best thing would be to throw an error message, but if that's not practical then it'd at least be nice to have the behavior be defined.

Changed files:
jax/random.py, tests/random_test.py

diff --git a/jax/random.py b/jax/random.py
index 8b00c9ca7..8b4af32d3 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -245,6 +245,10 @@ def randint(key, shape, minval, maxval, dtype=onp.int32):
   if nbits not in (32, 64):
     raise TypeError("randint only accepts 32- or 64-bit dtypes.")
 
+  # if we don't have minval < maxval, just always return minval
+  # https://github.com/google/jax/issues/222
+  maxval = lax.max(lax.add(minval, onp.array(1, dtype)), maxval)
+
   # This algorithm is biased whenever (maxval - minval) is not a power of 2.
   # We generate double the number of random bits required by the dtype so as to
   # reduce that bias.
diff --git a/tests/random_test.py b/tests/random_test.py
index 46963f076..9aada5a38 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -149,6 +149,10 @@ class LaxRandomTest(jtu.JaxTestCase):
     self.assertFalse(onp.all(perm1 == x))  # seems unlikely!
     self.assertTrue(onp.all(onp.sort(perm1) == x))
 
+  def testIssue222(self):
+    x = random.randint(random.PRNGKey(10003), (), 0, 0)
+    assert x == 0
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax
Issue: Failing lax_numpy_indexing_test.py tests on Python 3.7

`python tests/lax_numpy_indexing_test.py --num_generated_cases=1` on Python 3.7 (fresh Anaconda install on a macOS). Numpy version 1.15.4.

Changed files:
jax/numpy/lax_numpy.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index ef34e02f4..2f676424d 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1592,8 +1592,8 @@ def _rewriting_take(arr, idx, axis=0):
       axis += isinstance(elt, slice)   # advance axis index if not eliminated
     unexpanded_shape_itr = iter(result.shape)
     result_shape = tuple(1 if elt is None else next(unexpanded_shape_itr)
-                         for elt in canonical_idx if not isinstance(elt, int))
-    return lax.reshape(result, result_shape)
+                         for elt in canonical_idx if isinstance(elt, (type(None), slice)))
+    return lax.reshape(result, result_shape) if result_shape else result
 
   # Handle advanced indexing (non-tuple sequence, ndarray of dtype int or bool,
   # or a tuple with at least one sequence object).
Repository: google/jax
Issue: PYTHONOPTIMIZE=1 causes KeyErrors

### Description

When I set the flag: `export PYTHONOPTIMIZE=1`, the code:
```
import jax.numpy as jnp
m = jnp.eye(5)
a = jnp.ones(5)
jnp.linalg.solve(m,a)
```
returns the following error:
```
The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "error.py", line 6, in <module>
    jnp.linalg.solve(m,a)
  File "/home/ben/.../python3.8/site-packages/jax/_src/numpy/linalg.py", line 553, in solve
    return lax_linalg._solve(a, b)
KeyError: 139645559661888
```

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

jax v0.3.21, jaxlib v0.3.20

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

CPU

### Additional system info

Ubuntu, jax is installed into a site_packages repo for ROS.

### NVIDIA GPU info

_No response_

Changed files:
jax/interpreters/partial_eval.py

diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index bfacfe16b..78e71da9b 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -898,7 +898,8 @@ def tracers_to_jaxpr(
 
   def newvar(t: JaxprTracer) -> Var:
     var = gensym(type_substitute(t.aval))
-    assert t_to_var.setdefault(id(t), var) is var
+    var_ = t_to_var.setdefault(id(t), var)
+    assert var is var_
     return var
 
   def type_substitute(aval: AbstractValue) -> AbstractValue:
Repository: google/jax
Issue: BUG: can create 64-bit arrays using dlpack when x64 is disabled

### Description

Repro:
```python
import tensorflow as tf
import jax
import jax.dlpack

assert not jax.config.jax_enable_x64
x = jax.dlpack.from_dlpack(tf.experimental.dlpack.to_dlpack(tf.ones(5, tf.int64)))
assert x.dtype == jax.numpy.int64
```

This can cause loud and confusing errors down the line when our MHLO expects 32-bit numbers but gets 64-bit. For example, with our illegal 64-bit `x` array, try doing `x[0]`.

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

HEAD/0.3.15

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

CPU

### Additional System Info

_No response_

Changed files:
tests/array_interoperability_test.py, jax/_src/dlpack.py

diff --git a/jax/_src/dlpack.py b/jax/_src/dlpack.py
index b861370d4..64d79ee5a 100644
--- a/jax/_src/dlpack.py
+++ b/jax/_src/dlpack.py
@@ -70,4 +70,5 @@ def from_dlpack(dlpack):
   xla_shape = buf.xla_shape()
   assert not xla_shape.is_tuple()
   aval = core.ShapedArray(xla_shape.dimensions(), xla_shape.numpy_dtype())
-  return dispatch.maybe_create_array_from_da(buf, aval, buf.device())
+  return jnp.asarray(           # asarray ensures dtype canonicalization
+      dispatch.maybe_create_array_from_da(buf, aval, buf.device()))
diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py
index cc2f3a6c0..2da473275 100644
--- a/tests/array_interoperability_test.py
+++ b/tests/array_interoperability_test.py
@@ -137,6 +137,14 @@ class DLPackTest(jtu.JaxTestCase):
     y = tf.experimental.dlpack.from_dlpack(dlpack)
     self.assertAllClose(np, y.numpy())
 
+  @unittest.skipIf(not tf, "Test requires TensorFlow")
+  def testTensorFlowToJaxInt64(self):
+    # See https://github.com/google/jax/issues/11895
+    x = jax.dlpack.from_dlpack(
+        tf.experimental.dlpack.to_dlpack(tf.ones((2, 3), tf.int64)))
+    dtype_expected = jnp.int64 if config.x64_enabled else jnp.int32
+    self.assertEqual(x.dtype, dtype_expected)
+
   @jtu.sample_product(
     shape=all_shapes,
     dtype=torch_dtypes,
@@ -182,6 +190,14 @@ class DLPackTest(jtu.JaxTestCase):
     y = torch.utils.dlpack.from_dlpack(dlpack)
     self.assertAllClose(np, y.cpu().numpy())
 
+  @unittest.skipIf(not torch, "Test requires PyTorch")
+  def testTorchToJaxInt64(self):
+    # See https://github.com/google/jax/issues/11895
+    x = jax.dlpack.from_dlpack(
+        torch.utils.dlpack.to_dlpack(torch.ones((2, 3), dtype=torch.int64)))
+    dtype_expected = jnp.int64 if config.x64_enabled else jnp.int32
+    self.assertEqual(x.dtype, dtype_expected)
+
   @jtu.sample_product(
     shape=all_shapes,
     dtype=torch_dtypes,
Repository: google/jax
Issue: jax.numpy.linalg.lstsq throws IndexError with zero rows

### Description

[jax.numpy.linalg.lstsq](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.lstsq.html) throws an IndexError when the number of rows is zero:

```python3
import numpy as np
import jax.numpy as jnp

A = np.zeros([0, 3])
b = np.zeros([0])
x, residuals, rank, s = np.linalg.lstsq(A, b, None)
print(x) # [0. 0. 0.]

A = jnp.zeros([0, 3])
b = jnp.zeros([0])
x, residuals, rank, s = jnp.linalg.lstsq(A, b, None)
# IndexError: index is out of bounds for axis 0 with size 0
```

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

jax 0.3.21, jaxlib 0.3.20

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

CPU

### Additional system info

Python 3.10.7, macOS 11.7

### NVIDIA GPU 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 49cbbeb51..0c5f07df4 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -601,17 +601,22 @@ def _lstsq(a: ArrayLike, b: ArrayLike, rcond: Optional[float], *,
       f"{b.ndim}-dimensional array given. Array must be one or two-dimensional")
   m, n = a.shape
   dtype = a.dtype
-  if rcond is None:
-    rcond = jnp.finfo(dtype).eps * max(n, m)
+  if a.size == 0:
+    s = jnp.empty(0, dtype=a.dtype)
+    rank = jnp.array(0, dtype=int)
+    x = jnp.empty((n, *b.shape[1:]), dtype=a.dtype)
   else:
-    rcond = jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
-  u, s, vt = svd(a, full_matrices=False)
-  mask = s >= jnp.array(rcond, dtype=s.dtype) * s[0]
-  rank = mask.sum()
-  safe_s = jnp.where(mask, s, 1).astype(a.dtype)
-  s_inv = jnp.where(mask, 1 / safe_s, 0)[:, jnp.newaxis]
-  uTb = jnp.matmul(u.conj().T, b, precision=lax.Precision.HIGHEST)
-  x = jnp.matmul(vt.conj().T, s_inv * uTb, precision=lax.Precision.HIGHEST)
+    if rcond is None:
+      rcond = jnp.finfo(dtype).eps * max(n, m)
+    else:
+      rcond = jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
+    u, s, vt = svd(a, full_matrices=False)
+    mask = s >= jnp.array(rcond, dtype=s.dtype) * s[0]
+    rank = mask.sum()
+    safe_s = jnp.where(mask, s, 1).astype(a.dtype)
+    s_inv = jnp.where(mask, 1 / safe_s, 0)[:, jnp.newaxis]
+    uTb = jnp.matmul(u.conj().T, b, precision=lax.Precision.HIGHEST)
+    x = jnp.matmul(vt.conj().T, s_inv * uTb, precision=lax.Precision.HIGHEST)
   # Numpy returns empty residuals in some cases. To allow compilation, we
   # default to returning full residuals in all cases.
   if numpy_resid and (rank < n or m <= n):
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 0169380c0..7767f4619 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -890,6 +890,9 @@ class NumpyLinalgTest(jtu.JaxTestCase):
           ((4, 6), (4,)),
           ((6, 6), (6, 1)),
           ((8, 6), (8, 4)),
+          ((0, 3), (0,)),
+          ((3, 0), (3,)),
+          ((3, 1), (3, 0)),
       ]
     ],
     rcond=[-1, None, 0.5],
Repository: google/jax
Issue: jax.random.categorical doesn't respect shape parameter with broadcasted logits

### Description

The docs for `jax.random.categorical` say that the output of `jax.random.categorical` is `shape` when `shape` is provided. However, when provided `logits` which have to be broadcasted with `shape`, it just ignores the shape parameter after checking that it's broadcast compatible. I believe this can be fixed by broadcasting logits appropriately.

```
import jax
out = jax.random.categorical(jax.random.PRNGKey(0), logits=jax.numpy.ones( (4, 1, 10)), axis=-1, shape=(4, 8))

# out.shape is (4, 1), should be (4, 8)

assert out.shape == (4, 8), f"out.shape is {out.shape}, should be (4, 8)"
```



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

0.3.23

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

CPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

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

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 0e5e01dce..1292c8350 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1327,10 +1327,12 @@ def categorical(key: KeyArray,
     shape = tuple(shape)
     _check_shape("categorical", shape, batch_shape)
 
-  sample_shape = shape[:len(shape)-len(batch_shape)]
+  shape_prefix = shape[:len(shape)-len(batch_shape)]
+  logits_shape = list(shape[len(shape) - len(batch_shape):])
+  logits_shape.insert(axis % len(logits_arr.shape), logits_arr.shape[axis])
   return jnp.argmax(
-      gumbel(key, sample_shape + logits_arr.shape, logits_arr.dtype) +
-      lax.expand_dims(logits_arr, tuple(range(len(sample_shape)))),
+      gumbel(key, (*shape_prefix, *logits_shape), logits_arr.dtype) +
+      lax.expand_dims(logits_arr, tuple(range(len(shape_prefix)))),
       axis=axis)
 
 
diff --git a/tests/random_test.py b/tests/random_test.py
index 5cebbfd63..20802e990 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1429,6 +1429,19 @@ class LaxRandomTest(jtu.JaxTestCase):
     # just lower, don't run, takes too long
     jax.jit(f).lower()
 
+  @jtu.sample_product(shape=[(3, 4)],
+                      logits_shape_base=[(3, 4), (3, 1), (1, 4)],
+                      axis=[-3, -2, -1, 0, 1, 2])
+  def test_categorical_shape_argument(self, shape, logits_shape_base, axis):
+    # https://github.com/google/jax/issues/13124
+    logits_shape = list(logits_shape_base)
+    logits_shape.insert(axis % (len(logits_shape_base) + 1), 10)
+    assert logits_shape[axis] == 10
+    logits = jnp.ones(logits_shape)
+    samples = jax.random.categorical(jax.random.PRNGKey(0), logits=logits,
+                                     axis=axis, shape=shape)
+    self.assertEqual(samples.shape, shape)
+
 
 class KeyArrayTest(jtu.JaxTestCase):
   # Key arrays involve:
Repository: google/jax
Issue: jax.numpy.lcm int8 overflow difference with np

### Description

The following code in jax

```python
import jax.numpy as jnp
a = jnp.array([5], dtype=jnp.int8)
b = jnp.array([26], dtype=jnp.int8)
print(jnp.lcm(a, b))
```
produces `[126]` 

while in numpy (checked in 1.23.0 and 1.21.6) (and pytorch 1.11 and 1.12 ) it gives `[-126]` keeping true to the overflow for an `int8` dtype.

Is this difference of results between **`np`** and **`jnp`** desired here?

Surely, `jax.numpy` cannot follow NumPy API in every single output but this is something achievable in this case.
The cause is perhaps that jax uses `abs` at the end of lcm operation. I suppose it should suffice to apply `abs` on `x1` and `x2` as `gcd` does.

Note: Mathematically lcm(5,26)=130, so getting a negative value (as with numpy) can at least signal a warning but getting a positive (but incorrect) value close to the correct value can easily mislead.

int8 dtype was just picked for the sake of simpler example but I am guessing this won't be restricted to int8 dtype.

If you want then I can create a PR for this and in case I m totally mistakened please do enlighten me. Thank you very much.

### What jax/jaxlib version are you using?

jax 0.3.21, jaxlib v0.3.20. Also checked with jax master repo on GitHub.

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.8.10, Ubuntu 20, Google Colab

### NVIDIA GPU info

_No response_

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

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index c1d28b0db..e508f8ebd 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -4367,11 +4367,12 @@ def gcd(x1: ArrayLike, x2: ArrayLike) -> Array:
 def lcm(x1: ArrayLike, x2: ArrayLike) -> Array:
   _check_arraylike("lcm", x1, x2)
   x1, x2 = _promote_dtypes(x1, x2)
+  x1, x2 = abs(x1), abs(x2)
   if not issubdtype(_dtype(x1), integer):
     raise ValueError("Arguments to jax.numpy.lcm must be integers.")
   d = gcd(x1, x2)
   return where(d == 0, _lax_const(d, 0),
-               abs(multiply(x1, floor_divide(x2, d))))
+               multiply(x1, floor_divide(x2, d)))
 
 
 @_wraps(np.extract)
diff --git a/tests/lax_numpy_operators_test.py b/tests/lax_numpy_operators_test.py
index f852e6909..2d63f26fe 100644
--- a/tests/lax_numpy_operators_test.py
+++ b/tests/lax_numpy_operators_test.py
@@ -282,6 +282,7 @@ JAX_COMPOUND_OP_RECORDS = [
               all_shapes, jtu.rand_small_positive, []),
     op_record("gcd", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
     op_record("lcm", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
+    op_record("lcm", 2, [np.int8], all_shapes, jtu.rand_not_small, [])
 ]
 
 JAX_BITWISE_OP_RECORDS = [
Repository: google/jax
Issue: [sparse] bcoo_extract incorrect for zero sparse dimensions

Found in the course of #13624 

Minimal repro:
```python
import jax.numpy as jnp
from jax.experimental import sparse

mat = jnp.zeros(3, dtype='float32')
indices = jnp.zeros((3, 2, 0), dtype='int32')

jax.jit(sparse.bcoo_extract)(indices, mat)  # Error
```

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 247434523..81310b282 100644
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -400,7 +400,7 @@ def bcoo_extract(indices: Array, mat: Array) -> Array:
 @bcoo_extract_p.def_impl
 def _bcoo_extract_impl(indices, mat):
   mat = jnp.asarray(mat)
-  n_batch, n_sparse, _, _ = _validate_bcoo_indices(indices, mat.shape)
+  n_batch, n_sparse, _, nse = _validate_bcoo_indices(indices, mat.shape)
 
   ind_slices = tuple(np.zeros(s, int) if i_s == 1 else np.arange(s)
                      for s, i_s in zip(mat.shape[:n_batch], indices.shape[:n_batch]))
@@ -412,8 +412,13 @@ def _bcoo_extract_impl(indices, mat):
   batch_ind = tuple(grid)[:-1]
 
   if not sparse_ind + batch_ind:
-    return mat[None]
-  return mat.at[batch_ind + sparse_ind].get(mode='fill', fill_value=0)
+    result = mat[None]
+  else:
+    result = mat.at[batch_ind + sparse_ind].get(mode='fill', fill_value=0)
+  if n_sparse == 0 and nse != 1:
+    result = lax.broadcast_in_dim(
+      result, _tuple_replace(result.shape, n_batch, nse), range(result.ndim))
+  return result
 
 @bcoo_extract_p.def_abstract_eval
 def _bcoo_extract_abstract_eval(indices, mat):
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
index 7a9131082..27b029d0a 100644
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -904,6 +904,17 @@ class BCOOTest(sptu.SparseTestCase):
     self.assertEqual(j1.shape, data.shape + M.shape)
     self.assertEqual(hess.shape, data.shape + 2 * M.shape)
 
+  def test_bcoo_extract_zero_nse(self):
+    # Regression test for https://github.com/google/jax/issues/13653
+
+    # (n_batch, n_sparse, n_dense) = (1, 0, 0), nse = 2
+    args_maker = lambda: (jnp.zeros((3, 2, 0), dtype='int32'), jnp.arange(3))
+    self._CompileAndCheck(sparse.bcoo_extract, args_maker)
+
+    # (n_batch, n_sparse, n_dense) = (0, 0, 1), nse = 2
+    args_maker = lambda: (jnp.zeros((2, 0), dtype='int32'), jnp.arange(3))
+    self._CompileAndCheck(sparse.bcoo_extract, args_maker)
+
   @jtu.sample_product(
     [dict(shape=shape, n_batch=n_batch, n_dense=n_dense)
       for shape in [(), (5,), (5, 8), (3, 4, 5), (3, 4, 3, 2)]
Repository: google/jax
Issue: `replace` argument changes how `jax.random.choice` treats `shape`

Using jax version 0.1.75 on a colab cpu runtime.
```python
import jax

key = jax.random.PRNGKey(0)
jax.random.choice(key, 5, 2, replace=False)
jax.random.choice(key, 5, 2, replace=True)
```
Error:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-56696338bc13> in <module>()
      1 key = jax.random.PRNGKey(0)
      2 jax.random.choice(key, 5, 2, replace=False)
----> 3 jax.random.choice(key, 5, 2, replace=True)

2 frames
/usr/local/lib/python3.6/dist-packages/jax/random.py in choice(key, a, shape, replace, p)
    566   if p is None:
    567     if replace:
--> 568       ind = randint(key, shape, 0, n_inputs)
    569       result = ind if a.ndim == 0 else a[ind]
    570     else:

/usr/local/lib/python3.6/dist-packages/jax/random.py in randint(key, shape, minval, maxval, dtype)
    417   """
    418   dtype = dtypes.canonicalize_dtype(dtype)
--> 419   shape = abstract_arrays.canonicalize_shape(shape)
    420   return _randint(key, shape, minval, maxval, dtype)
    421 

/usr/local/lib/python3.6/dist-packages/jax/core.py in canonicalize_shape(shape)
   1078          "got {}.")
   1079   if any(isinstance(x, Tracer) and isinstance(get_aval(x), ShapedArray)
-> 1080          and not isinstance(get_aval(x), ConcreteArray) for x in shape):
   1081     msg += ("\\nIf using `jit`, try using `static_argnums` or applying `jit` to "
   1082             "smaller subfunctions.")

TypeError: 'int' object is not iterable
```

Changed files:
jax/random.py, tests/random_test.py

diff --git a/jax/random.py b/jax/random.py
index 985bccae4..2c775d924 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -556,6 +556,9 @@ def choice(key, a, shape=(), replace=True, p=None):
   Returns:
     An array of shape `shape` containing samples from `a`.
   """
+  if not isinstance(shape, Sequence):
+    raise TypeError("shape argument of jax.random.choice must be a sequence, "
+                    f"got {shape}")
   a = jnp.asarray(a)
   if a.ndim not in [0, 1]:
     raise ValueError("a must be an integer or 1-dimensional")
diff --git a/tests/random_test.py b/tests/random_test.py
index 7151955fe..e665bf0f2 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -848,6 +848,13 @@ class LaxRandomTest(jtu.JaxTestCase):
       self.assertAllClose(
           counts[1]/ num_samples, 0.5, rtol=1e-02, atol=1e-02)
 
+  def testChoiceShapeIsNotSequenceError(self):
+    key = random.PRNGKey(0)
+    with self.assertRaises(TypeError):
+      random.choice(key, 5, 2, replace=False)
+    with self.assertRaises(TypeError):
+      random.choice(key, 5, 2, replace=True)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax
Issue: Jit functions can't necessarily specialize on argument shapes

I'll disclaim that I'm not sure if this is a bug that should be fixed in Jax, or if the Jax "Sharp Bits" documentation should discuss this.

I originally stumbled upon an issue where my jit function was behaving strangely, and behaving differently than the non-jit version. Upon going over the Sharp Bits, I think this pertains to the "functions with argument-**value dependent shapes**" section. At first I thought it was simply my mistake, but then it dawned on me that the function doesn't in fact have value dependent shapes.

Here's a reproducible example:

```python
import jax
import jax.numpy as jnp

def fn(a, qs, idx):
  return a[idx, :, qs].mean(axis=-1).squeeze()

jit_fn = jax.jit(fn)

a = jnp.ones((24, 8, 11))
qs = jnp.arange(11)
fn(a, qs, 4).shape # outputs (8,) as expected
jit_fn(a, qs, 4).shape # outputs (11,)
``` 

Note that the shape of `fn` does not actually depend on the values of `qs`. In fact, none of the "intermediate steps" of the function produce any array whose shape depends on the values of `qs`. I fixed the problem by rewriting the functions as follows:

```python
def better_fn(a, idx):
  return a[idx, :].mean(axis=-1).squeeze()

jit_better_fn = jax.jit(better_fn)

b = a[:, :, qs]
better_fn(b, idx).shape # outputs (8,)
jit_better_fn(b, idx).shape # outputs (8,)
```

This is more evidence to me that the Sharp Bits section of the documentation applies to this case, however I believe that the statement that "specializing on argument shapes is ok" from the documentation is ambiguous.

Changed files:
jax/numpy/lax_numpy.py, tests/lax_numpy_indexing_test.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index cff076166..191fbc39b 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -3993,7 +3993,7 @@ def _index_to_gather(x_shape, idx):
     idx_no_nones = [(i, d) for i, d in enumerate(idx) if d is not None]
     advanced_pairs = (
       (asarray(e), i, j) for j, (i, e) in enumerate(idx_no_nones)
-      if isinstance(e, (Sequence, ndarray)))
+      if isscalar(e) or isinstance(e, (Sequence, ndarray)))
     advanced_pairs = ((_normalize_index(e, x_shape[j]), i, j)
                       for e, i, j in advanced_pairs)
     advanced_indexes, idx_advanced_axes, x_advanced_axes = zip(*advanced_pairs)
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index c5c732c6c..2b8d11dae 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -148,6 +148,11 @@ STATIC_INDEXING_TESTS = [
         IndexSpec(shape=(3,), indexer=()),
         IndexSpec(shape=(3, 4), indexer=()),
     ]),
+    ("TupleOfIntAndSliceAndIntArray", [
+        IndexSpec(shape=(3, 2, 3), indexer=(0, slice(None), np.arange(3))),
+        IndexSpec(shape=(3, 2, 3), indexer=(np.int32(1), slice(None), np.arange(3))),
+        IndexSpec(shape=(3, 2, 3), indexer=(np.array(2), slice(None), np.arange(3))),
+    ]),
 ]
 
 STATIC_INDEXING_GRAD_TESTS = [
Repository: google/jax
Issue: Missing check in shape rule of reduce_window

`reduce_window` expects to build a reduction function of type `T -> T -> T`, where `T` is the type of the operand and both arguments are scalar. If `init_value` is not a scalar, the reduction function that is constructed is invalid, and we get a `RuntimeError`.

This behavior can be reproduced using:
```python
from jax import lax
from jax import numpy as jnp

params = { "operand": jnp.ones((4, 4), dtype=jnp.int32)
         , "init_value": jnp.zeros((1,), dtype=jnp.int32) 
         , "computation": lax.max
         , "window_dimensions": (2, 2)
         , "window_strides": (2, 2)
         , "padding": "VALID"
         , "base_dilation": (1, 1)
         , "window_dilation": (1, 1)
         }
lax.reduce_window(**params)
```
which results in
```bash
RuntimeError: Invalid argument: Reduction function must return a scalar or tuple of scalars but returns shape: s32[1]: 
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```

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

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 3c030e52f..828704fb2 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -4882,6 +4882,10 @@ def _reduce_window_shape_rule(operand, init_value, *, jaxpr, consts,
     msg = ("reduce_window got inconsistent dtypes for operand and init_value: "
            " got operand dtype {} and init_value dtype {}.")
     raise TypeError(msg.format(operand.dtype, init_value.dtype))
+  if init_value.shape != ():
+    msg = ("reduce_window expected init_value to be a scalar but init_value "
+           "has shape {}.")
+    raise TypeError(msg.format(init_value.shape))
   return _common_reduce_window_shape_rule(
     operand, window_dimensions, window_strides, padding, base_dilation,
     window_dilation)
diff --git a/tests/lax_test.py b/tests/lax_test.py
index db6fc9c02..31ab3c82f 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2031,6 +2031,22 @@ class LaxTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(ValueError, msg):
       lax.conv_general_dilated(lhs, rhs, **kwargs)
 
+  def test_reduce_window_scalar_init_value_shape_rule(self):
+    # https://github.com/google/jax/issues/4574
+    args = { "operand": np.ones((4, 4), dtype=np.int32)
+           , "init_value": np.zeros((1,), dtype=np.int32)
+           , "computation": lax.max
+           , "window_dimensions": (2, 2)
+           , "window_strides": (2, 2)
+           , "padding": "VALID"
+           , "base_dilation": (1, 1)
+           , "window_dilation": (1, 1)
+           }
+
+    msg = (r"reduce_window expected init_value to be a scalar but init_value "
+           r"has shape \\(1,\\).")
+    with self.assertRaisesRegex(TypeError, msg):
+      lax.reduce_window(**args)
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):
Repository: google/jax
Issue: Inconsistent convention for complex derivatives

The convention for the handling of complex derivatives in JAX seems to be mixed. Consider the following two implementations of the abolute value function.
```py
import jax.numpy as np
from jax import grad

def f(x):
    return np.abs(x)

def g(x):
    return np.sqrt(np.real(x)**2 + np.imag(x)**2)

df = grad(f)
dg = grad(g)

print('gradient by df: ',df(1.0+0.5j))
print('gradient by dg: ',dg(1.0+0.5j))
```
The respective results differ by complex conjugation:
```
gradient by df:  (0.8944272-0.4472136j)
gradient by dg:  (0.8944272+0.4472136j)
```

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 f70cb0df2..7692a59d5 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2306,7 +2306,7 @@ real_p = unop(_complex_basetype, _complex, 'real')
 ad.deflinear(real_p, lambda t: [complex(t, np.zeros((), _dtype(t)))])
 
 imag_p = unop(_complex_basetype, _complex, 'imag')
-ad.defjvp(imag_p, lambda g, _: real(mul(_const(g, -1j), g)))
+ad.deflinear(imag_p, lambda t: [complex(np.zeros((), _dtype(t)), neg(t))])
 
 _complex_dtype = lambda dtype, *args: (np.zeros((), dtype) + np.zeros((), np.complex64)).dtype
 complex_p = naryop(_complex_dtype, [_complex_elem_types, _complex_elem_types],
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
index 4840a69dd..2c737bb14 100644
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -1067,6 +1067,20 @@ class LaxAutodiffTest(jtu.JaxTestCase):
     grad_fn = jax.grad(jax.grad(jax.grad(jax.grad(jax.grad(jax.grad(inv))))))
     self.assertAllClose(np.float32(0.0439453125), grad_fn(np.float32(4.)))
 
+  def test_linear_transpose_real(self):
+    f = lambda x: x.real
+    transpose = api.linear_transpose(f, 1.j)
+    actual, = transpose(1.)
+    expected = 1.
+    self.assertEqual(actual, expected)
+
+  def test_linear_transpose_imag(self):
+    f = lambda x: x.imag
+    transpose = api.linear_transpose(f, 1.j)
+    actual, = transpose(1.)
+    expected = -1.j
+    self.assertEqual(actual, expected)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax
Issue: repr(sharded_device_array) should print ShardedDeviceArray

Just noticed that we had a trivial regression in how repr(sharded_device_array) works:

```python
In [1]: from jax.api import device_put_sharded, devices

In [2]: device_put_sharded([1], devices())
Out[2]: DeviceArray([1], dtype=int32)
```

We use repr(x) to show the type of x, e.g. in demos, so we probably want to put this back.

Changed files:
jax/interpreters/xla.py, tests/pmap_test.py

diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 6a8cba5c8..4be8b3c22 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -1167,8 +1167,7 @@ for device_array in [_DeviceArray, _CppDeviceArray]:
 
   def __repr__(self):
     line_width = np.get_printoptions()["linewidth"]
-    # TODO(jblespia): Should we put back self.__class__.__name__ ?
-    prefix = '{}('.format("DeviceArray")
+    prefix = '{}('.format(self.__class__.__name__.lstrip('_'))
     s = np.array2string(self._value, prefix=prefix, suffix=',',
                         separator=', ', max_line_width=line_width)
     dtype_str = 'dtype={})'.format(self.dtype.name)
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
index 4315ad106..c1880fdce 100644
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -2091,6 +2091,10 @@ class ShardedDeviceArrayTest(jtu.JaxTestCase):
     self.assertTrue(all(b.device() == d for b, d in zip(y2.device_buffers, devices)))
     self.assertArraysEqual(y2, np.stack([xs['b'] for _ in devices]))
 
+  def test_repr(self):
+    x = jax.device_put_replicated(1, jax.devices())
+    self.assertStartsWith(repr(x), 'ShardedDeviceArray')
+
 
 class SpecToIndicesTest(jtu.JaxTestCase):