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

Issue:
psum over 32 bit numpy arrays returns 64 bit values when 64 bits are enabled
```python
import jax
import jax.lax as lax
import jax.numpy as jnp
import functools

jax.config.update('jax_enable_x64', True)

@functools.partial(jax.pmap, axis_name='i')
def foo(_):
  # Note that I never explicitly construct a numpy array, grad returns it sometimes.
  x = jax.grad(lambda x: jnp.array(0., jnp.float32))(jnp.array(0., jnp.float32))
  print(type(x))
  print(x.dtype)
  y = jax.lax.psum(x, axis_name='i')
  print(type(y))
  print(y.dtype)
  return y
foo(jnp.arange([jax.device_count()], dtype=jnp.float32))
```

Outputs:
```
<class 'numpy.ndarray'>
float32
<class 'numpy.float64'>
float64
```



Discussion:

Resolving patch:
diff --git a/jax/_src/lax/parallel.py b/jax/_src/lax/parallel.py
--- a/jax/_src/lax/parallel.py
+++ b/jax/_src/lax/parallel.py
@@ -759,7 +759,7 @@ def pos_reduce(x):
       size = len(axis_index_groups[0])
     else:
       size = prod([core.axis_frame(name).size for name in named_axes])  # type: ignore
-    return tuple(size * pos_reduce(x) for x in args)
+    return tuple(lax._const(x, size) * pos_reduce(x) for x in args)
   return core.Primitive.bind(
       psum_p, *args, axes=axes, axis_index_groups=axis_index_groups)
 


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1695,6 +1695,21 @@ def testArgAllReduce(self, shape, dtype, axis, collective, bulk_op):
     expected = bulk_op(x, axis=axis)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_dtype={}".format(
+          jtu.format_shape_dtype_string((), dtype)),
+       "dtype": dtype}
+      for dtype in [np.float32, np.int32]
+  )
+  def testPmapDtype(self, dtype):
+    # Regression test for https://github.com/google/jax/issues/6022
+    @partial(pmap, axis_name='i')
+    def func(_):
+      return jax.lax.psum(dtype(0), axis_name='i')
+    unused_arg = jnp.arange(xla_bridge.device_count())
+    out_dtype = func(unused_arg).dtype
+    self.assertEqual(out_dtype, dtype)
+
 
 class VmapOfPmapTest(jtu.JaxTestCase):
Repository: google/jax

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


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
@@ -1993,7 +1993,10 @@ def _reduction_init_val(a, init_val):
     sign, info = np.sign(init_val), iinfo(a_dtype)
     return np.array(info.min if sign < 0 else info.max, dtype=a_dtype)
 
-_cast_to_bool = partial(lax.convert_element_type, new_dtype=bool_)
+def _cast_to_bool(operand):
+  with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=np.ComplexWarning)
+    return lax.convert_element_type(operand, bool_)
 
 @_wraps(np.sum, skip_params=['out'])
 def sum(a, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype=None,


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
@@ -785,9 +785,7 @@ def np_fun(x):
       res = res if not is_bf16_nan_test else res.astype(jnp.bfloat16)
       return res
     np_fun = _promote_like_jnp(np_fun, inexact)
-    np_fun = jtu.ignore_warning(category=np.ComplexWarning)(np_fun)
     jnp_fun = lambda x: jnp_op(x, axis, keepdims=keepdims)
-    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
     args_maker = lambda: [rng(shape, dtype)]
     tol = {np.float16: 0.002}
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, tol=tol)
Repository: google/jax

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


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
@@ -1111,7 +1111,8 @@ def flip(m, axis: Optional[Union[int, Tuple[int, ...]]] = None):
   _check_arraylike("flip", m)
   if axis is None:
     return lax.rev(m, list(range(len(shape(m)))))
-  return lax.rev(m, [_canonicalize_axis(axis, ndim(m))])
+  axis = _ensure_index_tuple(axis)
+  return lax.rev(m, [_canonicalize_axis(ax, ndim(m)) for ax in axis])
 
 
 @_wraps(np.fliplr)


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
@@ -3449,7 +3449,7 @@ def testTracingPrimitiveWithNoTranslationErrorMessage(self):
        "shape": shape, "dtype": dtype, "axis": axis}
       for shape in [(3,), (2, 3)]
       for dtype in default_dtypes
-      for axis in list(range(-len(shape), len(shape))) + [None]  # Test negative axes
+      for axis in list(range(-len(shape), len(shape))) + [None] + [tuple(range(len(shape)))]  # Test negative axes and tuples
     ))
   def testFlip(self, shape, dtype, axis):
     rng = jtu.rand_default(self.rng())
Repository: google/jax

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


Discussion:

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
@@ -5573,6 +5573,8 @@ def _select_and_scatter_add_transpose(
     t, source, operand, *, select_prim, window_dimensions, window_strides,
     padding):
   assert ad.is_undefined_primal(source) and not ad.is_undefined_primal(operand)
+  if type(t) is ad_util.Zero:
+    return [ad_util.Zero(source.aval), None]
   ones = (1,) * len(window_dimensions)
   source_t = _select_and_gather_add(t, operand, select_prim, window_dimensions,
                                     window_strides, padding, ones, ones)


Test patch:
Repository: google/jax

Issue:
jnp.quantile should return NaN for arrays containing NaNs
```python
import numpy as np
import jax.numpy as jnp

x = np.array([1, 2, np.nan])

print(np.quantile(x, 0.5))
# nan

print(jnp.quantile(x, 0.5))
# 2.0
```


Discussion:
Related to #6439

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
@@ -5228,6 +5228,7 @@ def _quantile(a, q, axis, interpolation, keepdims, squash_nans):
     index[axis] = high
     high_value = a[tuple(index)]
   else:
+    a = where(any(isnan(a), axis=axis, keepdims=True), nan, a)
     a = lax.sort(a, dimension=axis)
     n = a_shape[axis]
     q = lax.mul(q, _constant_like(q, n - 1))


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
@@ -4037,7 +4037,7 @@ def args_maker(): return []
              jtu.format_shape_dtype_string(a_shape, a_dtype),
              jtu.format_shape_dtype_string(q_shape, q_dtype),
              axis, keepdims, interpolation),
-         "a_rng": jtu.rand_some_nan if 'nan' in op else jtu.rand_default,
+         "a_rng": jtu.rand_some_nan,
          "q_rng": q_rng, "op": op,
          "a_shape": a_shape, "a_dtype": a_dtype,
          "q_shape": q_shape, "q_dtype": q_dtype, "axis": axis,
@@ -4068,6 +4068,9 @@ def testQuantile(self, op, a_rng, q_rng, a_shape, a_dtype, q_shape, q_dtype,
       args_maker = lambda: [a_rng(a_shape, a_dtype)]
     else:
       args_maker = lambda: [a_rng(a_shape, a_dtype), q_rng(q_shape, q_dtype)]
+
+    # TODO(jakevdp): remove this ignore_warning when minimum numpy version is 1.17.0
+    @jtu.ignore_warning(category=RuntimeWarning, message="Invalid value encountered.*")
     def np_fun(*args):
       args = [x if jnp.result_type(x) != jnp.bfloat16 else
               np.asarray(x, np.float32) for x in args]
Repository: google/jax

Issue:
jnp.nanmax does not work with int inputs
Just flagging that `jnp.nanmax` has issues when the input is int. For example:
```
#Numpy version:
np.nanmax(np.arange(3*3*3).reshape(3,3,3), axis=(1,2))
>> array([ 8, 17, 26])
```
versus:
```
jnp.nanmax(jnp.arange(3*3*3).reshape(3,3,3), axis=(1,2))
>> ValueError: cannot convert float NaN to integer
```
Converting to float works, as does jnp.nanmean. Issue seems to be converting nan to int here: 
https://github.com/google/jax/blob/52a41311c5666d0188c9f6966e3ef7e31ccb4e81/jax/numpy/lax_numpy.py#L275


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
@@ -2286,6 +2286,9 @@ def flatnonzero(a):
 def _nan_reduction(a, name, jnp_reduction, init_val, nan_if_all_nan,
                    axis=None, keepdims=None, **kwargs):
   _check_arraylike(name, a)
+  if not issubdtype(_dtype(a), inexact):
+    return jnp_reduction(a, axis=axis, keepdims=keepdims, **kwargs)
+
   out = jnp_reduction(where(isnan(a), _reduction_init_val(a, init_val), a),
                       axis=axis, keepdims=keepdims, **kwargs)
   if nan_if_all_nan:


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
@@ -339,7 +339,7 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     op_record("sum", 1, all_dtypes, all_shapes, jtu.rand_default, []),
     op_record("nanmean", 1, inexact_dtypes, nonempty_shapes, jtu.rand_some_nan,
               [], inexact=True),
-    op_record("nanprod", 1, inexact_dtypes, all_shapes, jtu.rand_some_nan, []),
+    op_record("nanprod", 1, all_dtypes, all_shapes, jtu.rand_some_nan, []),
     op_record("nansum", 1, number_dtypes, all_shapes, jtu.rand_some_nan, []),
 ]
 
@@ -370,8 +370,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
               inexact=True),
     op_record("std", 1, all_dtypes, nonempty_shapes, jtu.rand_default, [],
               inexact=True),
-    op_record("nanmax", 1, inexact_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
-    op_record("nanmin", 1, inexact_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
+    op_record("nanmax", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
+    op_record("nanmin", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
     op_record("nanvar", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan,
               [], inexact=True),
     op_record("nanstd", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan,
Repository: google/jax

Issue:
Issue better error message when jax.numpy.pad() pad_widths mismatches array rank
The following JAX code:

```
    y = jnp.pad(jnp.zeros((2,2)), [[0,0],[0,0],[0,0]], mode='wrap')
```
given an error:

```
ValueError: pad_width given unexpected structure: [[0 0]
 [0 0]
 [0 0]]. See docstring for valid pad_width formats.
```
This is because the length of the pad_widths list needs to match the rank of the array. Can we make this error message more explicit?


Discussion:
It would certainly be possible to be more explicit, but in this case it's tricky because there are so many possible reasons that the input is invalid. Here is the chain of options: https://github.com/google/jax/blob/9c63c772a97623453aecc0680dd4601cc8aa107a/jax/_src/numpy/lax_numpy.py#L2639-L2656

It's tough to guess what the user intended when there are so many possible correct inputs. Maybe in this case we could raise a more specific error if `nvals.ndim == 2 and nvals.shape[1] == 2 and nvals.shape[0] not in (1, nd)`. What do you think?
One simple thing might be just to report `np.ndm(x)` in the error message.

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
@@ -2658,8 +2658,8 @@ def _broadcast_to_pairs(nvals, nd, name):
     # pad
     return tuple((nvals.flat[0], nvals.flat[0]) for i in range(nd))
   else:
-    raise ValueError(f"{name} given unexpected structure: {nvals}. "
-                     f"See docstring for valid {name} formats.")
+    raise ValueError(f"jnp.pad: {name} with nd={nd} has unsupported shape {nvals.shape}. "
+                     f"Valid shapes are ({nd}, 2), (1, 2), (2,), (1,), or ().")
 
 
 @partial(jit, static_argnums=(1, 2, 4, 5, 6))


Test patch:
Repository: google/jax

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


Discussion:
The issue is actually that `argmax` is handled differently on GPU, and its implementation requires an equality comparison, which fails for NaNs.

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
@@ -5382,7 +5382,8 @@ def _argminmax_translation_rule(value_comparator, identity,
   x_index = xb.parameter(subc, 1, index_shape)
   y_value = xb.parameter(subc, 2, value_shape)
   y_index = xb.parameter(subc, 3, index_shape)
-  which_value = value_comparator(x_value, y_value)
+  which_value = xops.Or(value_comparator(x_value, y_value),
+                        xops.Ne(x_value, x_value))
   which_index = xops.Or(which_value, xops.And(xops.Eq(x_value, y_value),
                                               xops.Lt(x_index, y_index)))
   xops.Tuple(subc, [xops.Select(which_value, x_value, y_value),
@@ -5402,8 +5403,8 @@ def _argminmax_gpu_translation_rule(op, a, *, axes, index_dtype):
   idxs = tie_in(a, broadcasted_iota(index_dtype, a.shape, axis))
   maxval = np.array(dtypes.iinfo(index_dtype).max, dtype=index_dtype)
   maxval = broadcast(tie_in(a, maxval), a.shape)
-  mask_idxs = select(eq(a, expand_dims(op(a, (axis,)), (axis,))), idxs,
-                     maxval)
+  maxvals = expand_dims(op(a, (axis,)), (axis,))
+  mask_idxs = select(eq(a, maxvals) | ne(a, a), idxs, maxval)
   return _reduce_min(mask_idxs, (axis,))
 
 _argmin_translation_rule = partial(_argminmax_translation_rule, xops.Lt,


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2560,6 +2560,10 @@ def testArgMinMaxWeakType(self, jax_fn, weak_type):
     x_out_jit = api.jit(op)(x_in)
     self.assertEqual(dtypes.is_weakly_typed(x_out_jit), False)
 
+  def testArgMaxOfNanChoosesNaN(self):
+    self.assertEqual(lax.argmax(np.array([0., np.nan]), axis=0,
+                                index_dtype=np.int32), 1)
+
   @parameterized.named_parameters(jtu.cases_from_list(
         {"testcase_name": "_{}".format(rec.op),
          "op_name": rec.op, "rec_dtypes": rec.dtypes}
Repository: google/jax

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



Discussion:
You're correct, although the function appears to be unused.

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


Test patch:
Repository: google/jax

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


Discussion:
Hi! I've got some time. I'd like to take up this issue. 
Awesome! Please take a look at the [contributing](https://jax.readthedocs.io/en/latest/contributing.html#contributing-code-using-pull-requests) doc, and please feel free to reach out if you have questions.

You can open a draft pull request if that would be helpful :grin:

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
@@ -1805,7 +1805,8 @@ def _split(op, ary, indices_or_sections, axis=0):
   size = ary.shape[axis]
   if isinstance(indices_or_sections, (tuple, list) + _arraylike_types):
     indices_or_sections = np.array(
-        [core.concrete_or_error(np.int64, i_s, f"in jax.numpy.{op} argument 1")
+        [_canonicalize_axis(core.concrete_or_error(np.int64, i_s,
+                                                   f"in jax.numpy.{op} argument 1"), size)
          for i_s in indices_or_sections], np.int64)
     split_indices = np.concatenate([[np.int64(0)], indices_or_sections,
                                     [np.int64(size)]])


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
@@ -2925,8 +2925,8 @@ def testSliceWeakTypes(self, shape, dtype, weak_type, slc):
        "shape": shape, "num_sections": num_sections, "axis": axis,
        "dtype": dtype}
       for shape, axis, num_sections in [
-          ((3,), 0, 3), ((12,), 0, 3), ((12, 4), 0, 4), ((12, 4), 1, 2),
-          ((2, 3, 4), -1, 2), ((2, 3, 4), -2, 3)]
+          ((3,), 0, 3), ((12,), 0, 3), ((12, 4), 0, 4), ((3,), 0, [-1]),
+          ((12, 4), 1, 2), ((2, 3, 4), -1, 2), ((2, 3, 4), -2, 3)]
       for dtype in default_dtypes))
   def testSplitStaticInt(self, shape, num_sections, axis, dtype):
     rng = jtu.rand_default(self.rng())
Repository: google/jax

Issue:
Possible bug in shape mismatch error message
https://github.com/google/jax/blob/ecab743e5c4a4d87a6a6031cc2cbc599bb2e3383/jax/_src/lax/control_flow.py#L2513

Shouldn't it be elem.shape?


Discussion:
Definitely should be `elem.shape`. Thanks for letting us know!

Resolving patch:
diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -2510,7 +2510,7 @@ def combine(a_flat, b_flat):
   if not all(int(elem.shape[axis]) == num_elems for elem in elems_flat[1:]):
     raise ValueError('Array inputs to associative_scan must have the same '
                      'first dimension. (saw: {})'
-                     .format([elems.shape for elem in elems_flat]))
+                     .format([elem.shape for elem in elems_flat]))
 
 
   # Summary of algorithm:


Test patch:
Repository: google/jax

Issue:
`jnp.sinc` causes error with NaNs debugging enabled
The [current implementation of `sinc`](https://github.com/google/jax/blob/d39261497c12488632cf7c5d5f19d2c074a565dd/jax/_src/numpy/lax_numpy.py#L1061) function generates NaNs that are detected when `jax_debug_nans` is activated.

**To reproduce:**
```python
import jax
import jax.numpy as jnp
from jax.config import config 
config.update("jax_debug_nans", True)

x = jnp.array([0.])
x = jnp.sinc(x)
```

Output:
```python
---------------------------------------------------------------------------

FloatingPointError                        Traceback (most recent call last)

<ipython-input-5-ca7132c5c117> in <module>()
      6 x = jnp.array([0.])
      7 
----> 8 x = jnp.sinc(x)

[...]

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
    374   if dtypes.issubdtype(xla_shape.element_type(), np.inexact):
    375     if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 376       raise FloatingPointError(f"invalid value (nan) encountered in {name}")
    377     if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):
    378       raise FloatingPointError(f"invalid value (inf) encountered in {name}")

FloatingPointError: invalid value (nan) encountered in div
```


The error comes from the second branch of the `jnp.where` function [here](https://github.com/google/jax/blob/d39261497c12488632cf7c5d5f19d2c074a565dd/jax/_src/numpy/lax_numpy.py#L1068), which has NaNs where `safe_pi_x` is `0`. Those NaNs are filtered out by the `where` function but are nevertheless produced and detected by the NaNs debugger.



Discussion:
Thanks for the report - this should be fixed by #6940

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
@@ -1064,7 +1064,7 @@ def sinc(x):
   x, = _promote_dtypes_inexact(x)
   eq_zero = lax.eq(x, lax._const(x, 0))
   pi_x = lax.mul(lax._const(x, pi), x)
-  safe_pi_x = where(eq_zero, lax._const(x, 0), pi_x)
+  safe_pi_x = where(eq_zero, lax._const(x, 1), pi_x)
   return where(eq_zero, _sinc_maclaurin(0, pi_x),
                lax.div(lax.sin(safe_pi_x), safe_pi_x))
 


Test patch:
diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
--- a/tests/debug_nans_test.py
+++ b/tests/debug_nans_test.py
@@ -37,6 +37,10 @@ def setUp(self):
   def tearDown(self):
     config.update("jax_debug_nans", self.cfg)
 
+  def testSinc(self):
+    # Regression test for #6936
+    self.assertEqual(jnp.sinc(0.0), 1.0)
+
   def testSingleResultPrimitiveNoNaN(self):
     A = jnp.array([[1., 2.], [2., 3.]])
     ans = jnp.tanh(A)