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

Issue:
Rename jax.nn.swish to jax.nn.silu to give appropriate credit
The swish was originally coined the "SiLU" in https://arxiv.org/pdf/1606.08415.pdf and https://arxiv.org/abs/1702.03118 long before the swish paper. Renaming other peoples' exact same ideas is unacceptable and tensorflow's naming convention implicitly erases the research and work of people outside of Google.
This request inspired by a [recent discussion](https://www.reddit.com/r/MachineLearning/comments/hkiyir/r_google_has_a_credit_assignment_problem_in/) and a recent [tensorflow issue](https://github.com/tensorflow/tensorflow/issues/41066), but this problem has been brought up every few months for the past few years. In light of recent efforts to make the ML community more equitable and _fair_, this is a no-brainer and long overdue.

**Will this change the current api? How?**
jax.nn.swish will eventually be deprecated (jax is still new) and jax.nn.silu will be added and both of the aforementioned papers will be cited in the documentation.


Discussion:

Resolving patch:
diff --git a/jax/nn/__init__.py b/jax/nn/__init__.py
--- a/jax/nn/__init__.py
+++ b/jax/nn/__init__.py
@@ -22,6 +22,7 @@
   gelu,
   glu,
   hard_sigmoid,
+  hard_silu,
   hard_swish,
   hard_tanh,
   leaky_relu,
@@ -36,5 +37,6 @@
   soft_sign,
   softmax,
   softplus,
+  silu,
   swish,
 )
diff --git a/jax/nn/functions.py b/jax/nn/functions.py
--- a/jax/nn/functions.py
+++ b/jax/nn/functions.py
@@ -68,16 +68,18 @@ def sigmoid(x):
   """
   return expit(x)
 
-def swish(x):
-  r"""Swish activation function.
+def silu(x):
+  r"""SiLU activation function.
 
   Computes the element-wise function:
 
   .. math::
-    \mathrm{swish}(x) = x \cdot \mathrm{sigmoid}(x) = \frac{x}{1 + e^{-x}}
+    \mathrm{silu}(x) = x \cdot \mathrm{sigmoid}(x) = \frac{x}{1 + e^{-x}}
   """
   return x * sigmoid(x)
 
+swish = silu
+
 def log_sigmoid(x):
   r"""Log-sigmoid activation function.
 
@@ -292,12 +294,14 @@ def hard_sigmoid(x):
   """
   return relu6(x + 3.) / 6.
 
-def hard_swish(x):
-  r"""Hard Swish activation function
+def hard_silu(x):
+  r"""Hard SiLU activation function
 
   Computes the element-wise function
 
   .. math::
-    \mathrm{hard\_swish}(x) = x \cdot \mathrm{hard\_sigmoid}(x)
+    \mathrm{hard\_silu}(x) = x \cdot \mathrm{hard\_sigmoid}(x)
   """
   return x * hard_sigmoid(x)
+
+hard_swish = hard_silu


Test patch:
Repository: google/jax

Issue:
jax.numpy.take() fails on empty arrays (and np.take does not)
```
>>> import numpy as onp
>>> import jax.numpy as jnp
>>> onp.take([], [])
array([], dtype=float64)
>>> jnp.take([], [])
ValueError: start_indices must have an integer type
```
This error is caused because jax converts the second array to an empty float array. But another error is raised even if you work around the first one:
```
>>> jnp.take([], jnp.array([], dtype=int))
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.:
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```



Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -3334,7 +3334,7 @@ def take(a, indices, axis=None, out=None, mode=None):
 
   index_dims = len(shape(indices))
   slice_sizes = list(shape(a))
-  slice_sizes[axis] = 1
+  slice_sizes[axis] = _min(indices.size, 1)
   dnums = lax.GatherDimensionNumbers(
     offset_dims=tuple(
       list(range(axis)) +


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
@@ -2761,6 +2761,12 @@ def args_maker():
     self._CheckAgainstNumpy(jnp_op, np_op, args_maker)
     self._CompileAndCheck(jnp_op, args_maker)
 
+  def testTakeEmpty(self):
+    np.testing.assert_array_equal(
+      jnp.array([], dtype=jnp.float32),
+      jnp.take(jnp.array([], jnp.float32), jnp.array([], jnp.int32)))
+
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_ishape={}_axis={}".format(
           jtu.format_shape_dtype_string(x_shape, dtype), i_shape, axis),
Repository: google/jax

Issue:
jax.scipy.special zeta sometimes returns NaN
For large values of z, `jax.scipy.special.zeta` returns NaN when it should instead return 0.:

`jax.scipy.special.zeta(1e5, q=1.)`  # Return NaN.

Note: I suspect TF and Scipy have this same issue:

`sp.zeta(1e19, 40.)`   # Returns NaN in Scipy. I have to use a large value because Scipy is using internal float64 calculations.

`tf.math.zeta(1e10, q=30.)`  # Returns NaN in TF


Discussion:
@fehiepsi contributed this function originally, perhaps they have thoughts.
I think you can fix the issue by clipping [T1 term](https://github.com/google/jax/blob/master/jax/scipy/special.py#L208) by `jnp.clip(T1, a_max=jnp.finfo(dtype).max)`. Another way is to add a safemul when multiplying `0` and `inf` at the computation of [T](https://github.com/google/jax/blob/master/jax/scipy/special.py#L211), kind of
```
T = jnp.where(T0 == 0, 0., T0 * (dtype(0.5) + T1.sum(-1)))
```
That way, we don't face `NaN` issue, even for examples that failed on `scipy` and `tf`.

Resolving patch:
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -206,6 +206,7 @@ def zeta(x, q=None):
   T0 = (a + N) ** -s
   s_over_a = (s_ + np.arange(2 * M, dtype=M.dtype)) / (a_ + N)
   T1 = jnp.cumprod(s_over_a, -1)[..., ::2]
+  T1 = jnp.clip(T1, a_max=jnp.finfo(dtype).max)
   coefs = np.array(_BERNOULLI_COEFS[:T1.shape[-1]], dtype=dtype)
   T1 = T1 / coefs
   T = T0 * (dtype(0.5) + T1.sum(-1))


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -196,6 +196,11 @@ def testIssue980(self):
     self.assertAllClose(np.zeros((4,), dtype=np.float32),
                         lsp_special.expit(x))
 
+  def testIssue3758(self):
+    x = np.array([1e5, 1e19, 1e10], dtype=np.float32)
+    q = np.array([1., 40., 30.], dtype=np.float32)
+    self.assertAllClose(np.array([1., 0., 0.], dtype=np.float32), lsp_special.zeta(x, q))
+
   def testXlogyShouldReturnZero(self):
     self.assertAllClose(lsp_special.xlogy(0., 0.), 0., check_dtypes=False)
Repository: google/jax

Issue:
vmap mismatched dim exception has another exception composing the error msg
```python
import jax.numpy as np
def recursion(depth, seed, xt, xl):
  seed, sample_seed = jax.random.split(seed)
  if depth == 2:
    return xl
  xl = xl + tf.random.stateless_normal([], seed=sample_seed)
  return recursion(depth + 1, seed, xt, xl)
jax.vmap(recursion, in_axes=(None, None, 0, 0))(
    0, jax.random.PRNGKey(2), np.linspace(0, 1, 7), np.zeros([2]))
```

want: (something like what you get with `np.array(0)` for the first arg)
```
ValueError: vmap got inconsistent sizes for array axes to be mapped:
arg 0 has shape () and axis None is to be mapped
arg 1 has shape (2,) and axis None is to be mapped
arg 2 has shape (7,) and axis 0 is to be mapped
arg 3 has shape (2,) and axis 0 is to be mapped
so
arg 2 has an axis to be mapped of size 7
arg 3 has an axis to be mapped of size 2
```
but get
```
jax/api.py in <listcomp>(.0)
    908     if tree == tree_flatten((core.unit,) * tree.num_leaves)[1]:
    909       lines1 = ["arg {} has shape {} and axis {} is to be mapped"
--> 910                 .format(i, x.shape, d) for i, (x, d) in enumerate(zip(vals, dims))]
    911       sizes = collections.defaultdict(list)
    912       for i, (x, d) in enumerate(zip(vals, dims)):

AttributeError: 'int' object has no attribute 'shape'
```


Discussion:
Thanks for reporting this!

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -907,7 +907,7 @@ def _mapped_axis_size(tree, vals, dims, name):
     # TODO(mattjj,phawkins): add a way to inspect pytree kind more directly
     if tree == tree_flatten((core.unit,) * tree.num_leaves)[1]:
       lines1 = ["arg {} has shape {} and axis {} is to be mapped"
-                .format(i, x.shape, d) for i, (x, d) in enumerate(zip(vals, dims))]
+                .format(i, np.shape(x), d) for i, (x, d) in enumerate(zip(vals, dims))]
       sizes = collections.defaultdict(list)
       for i, (x, d) in enumerate(zip(vals, dims)):
         if d is not None:
@@ -919,11 +919,11 @@ def _mapped_axis_size(tree, vals, dims, name):
                   "axes" if len(idxs) > 1 else "an axis",
                   size)
                 for size, idxs in sizes.items()]
-      raise ValueError(msg.format("\n".join(lines1 + ["so"] + lines2))) from e
+      raise ValueError(msg.format("\n".join(lines1 + ["so"] + lines2))) from None
     else:
       sizes = [x.shape[d] if d is not None else None for x, d in zip(vals, dims)]
       sizes = tree_unflatten(tree, sizes)
-      raise ValueError(msg.format("the tree of axis sizes is:\n{}".format(sizes))) from e
+      raise ValueError(msg.format("the tree of axis sizes is:\n{}".format(sizes))) from None
 
 def pmap(fun: Callable, axis_name: Optional[AxisName] = None, *, in_axes=0,
          static_broadcasted_argnums: Union[int, Iterable[int]] = (),


Test patch:
Repository: google/jax

Issue:
jax.numpy.tile inconsistent with np.tile (when 0 in repeats)
```python
>>> import numpy as np
>>> import jax.numpy as jnp
>>> np.tile(np.array([0, 1, 2]), (1, 1, 2))
array([[[0, 1, 2, 0, 1, 2]]])
>>> jnp.tile(jnp.array([0, 1, 2]), (1, 2, 2))
DeviceArray([[[0, 1, 2, 0, 1, 2]]], dtype=int64)
>>> np.tile(np.array([0, 1, 2]), (1, 0, 2))
array([], shape(1, 0, 6), dtype=int64)
>>> jnp.tile(jnp.array([0, 1, 2]), (1, 0, 2)
...
ValueError: Need at least one array to concatenate
```

A fix locally is to wrap the `jax.numpy.tile` like so:

```python
        if 0 in repeats:
            return jnp.array([]).reshape(np.array(tensor_in.shape) * np.array(repeats))
        return jnp.tile(tensor_in, repeats)
```

jax.numpy.tile inconsistent with np.tile (when 0 in repeats)
```python
>>> import numpy as np
>>> import jax.numpy as jnp
>>> np.tile(np.array([0, 1, 2]), (1, 1, 2))
array([[[0, 1, 2, 0, 1, 2]]])
>>> jnp.tile(jnp.array([0, 1, 2]), (1, 2, 2))
DeviceArray([[[0, 1, 2, 0, 1, 2]]], dtype=int64)
>>> np.tile(np.array([0, 1, 2]), (1, 0, 2))
array([], shape(1, 0, 6), dtype=int64)
>>> jnp.tile(jnp.array([0, 1, 2]), (1, 0, 2)
...
ValueError: Need at least one array to concatenate
```

A fix locally is to wrap the `jax.numpy.tile` like so:

```python
        if 0 in repeats:
            return jnp.array([]).reshape(np.array(tensor_in.shape) * np.array(repeats))
        return jnp.tile(tensor_in, repeats)
```



Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2108,7 +2108,10 @@ def tile(A, reps):
   A = reshape(A, (1,) * (len(reps) - ndim(A)) + shape(A))
   reps = (1,) * (ndim(A) - len(reps)) + tuple(reps)
   for i, rep in enumerate(reps):
-    A = concatenate([A] * int(rep), axis=i)
+    if rep == 0:
+      A = A[tuple(slice(0 if j == i else None) for j in range(A.ndim))]
+    elif rep != 1:
+      A = concatenate([A] * int(rep), axis=i)
   return A
 
 @_wraps(np.concatenate)


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
@@ -1137,7 +1137,7 @@ def args_maker():
           jtu.format_shape_dtype_string(shape, dtype), reps),
        "shape": shape, "dtype": dtype, "reps": reps,
        "rng_factory": jtu.rand_default}
-      for reps in [(), (2,), (3, 4), (2, 3, 4)]
+      for reps in [(), (2,), (3, 4), (2, 3, 4), (1, 0, 2)]
       for shape, dtype in _shape_and_dtypes(all_shapes, default_dtypes)
       ))
   def testTile(self, shape, dtype, reps, rng_factory):
Repository: google/jax

Issue:
JAX arrays should implement nbytes
`nbytes` is a convenient shortcut for calculating the size of a NumPy arrays. The implement is simply `array.size * array.dtype.itemsize`.

But currently, we see:
```
>>> import jax.numpy as jnp
>>> jnp.zeros((1000, 1000)).nbytes
AttributeError: 'DeviceArray' object has no attribute 'nbytes'
```




Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -4377,6 +4377,11 @@ def _astype(arr, dtype):
   lax._check_user_dtype_supported(dtype, "astype")
   return lax.convert_element_type(arr, dtype)
 
+
+def _nbytes(arr):
+  return size(arr) * _dtype(arr).itemsize
+
+
 def _view(arr, dtype=None, type=None):
   if type is not None:
     raise NotImplementedError("`type` argument of array.view()")
@@ -4533,6 +4538,7 @@ def _operator_round(number, ndigits=None):
 setattr(ShapedArray, "imag", core.aval_property(imag))
 setattr(ShapedArray, "astype", core.aval_method(_astype))
 setattr(ShapedArray, "view", core.aval_method(_view))
+setattr(ShapedArray, "nbytes", core.aval_property(_nbytes))
 
 
 # Forward operators, methods, and properties on DeviceArray to lax_numpy
@@ -4548,6 +4554,7 @@ def _operator_round(number, ndigits=None):
 setattr(DeviceArray, "imag", property(imag))
 setattr(DeviceArray, "astype", _astype)
 setattr(DeviceArray, "view", _view)
+setattr(DeviceArray, "nbytes", property(_nbytes))
 
 
 # Extra methods that are handy


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
@@ -2662,6 +2662,20 @@ def testAstype(self):
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
     self._CompileAndCheck(jnp_op, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_{}".format(
+          jtu.format_shape_dtype_string(shape, dtype)),
+      "shape": shape, "dtype": dtype}
+      for shape in array_shapes
+      for dtype in all_dtypes))
+  def testNbytes(self, shape, dtype):
+    rng = jtu.rand_default(self.rng())
+    np_op = lambda x: np.asarray(x).nbytes
+    jnp_op = lambda x: jnp.asarray(x).nbytes
+    args_maker = lambda: [rng(shape, dtype)]
+    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
+    self._CompileAndCheck(jnp_op, args_maker)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_dtype={}".format(
           jtu.format_shape_dtype_string(shape, a_dtype), dtype),
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
```


Discussion:

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


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -848,6 +848,13 @@ def testRadamacher(self):
       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:
Nested pmap issue on GPU
Running the code with nested pmap/vmap in the reproducer below works with CPU/TPU but fails on GPU. The explicitly specified `cpu` backend when running `pmap` seems to be ignored, as there are clearly enough devices available.

Changing 

`updates_at_idxs = updates[..., None] * rlax.one_hot(indices, num_classes)` 

to 

`updates_at_idxs = updates[..., None]` 

makes things work on GPU as well. Not entirely sure how this is related to the mapped axes though.

- [Colab reproducer](https://colab.research.google.com/drive/1C8i4xUbu8wDWJKy0EvGhu2Vn7ptAky6W?usp=sharing)


Discussion:
I've further isolated the error and included it as an alternative reproducer in the provided colab. This fails on GPU:

```python
import chex
import jax
from jax import numpy as jnp

chex.set_n_cpu_devices(4)

def fn(indices):
  return jnp.equal(indices, jnp.arange(3)).astype(jnp.float32)

mapped_fn = jax.pmap(fn, axis_name='i', backend='cpu')
mapped_fn = jax.pmap(mapped_fn, axis_name='j', backend='cpu')

indices = np.array([[[2], [1]], [[0], [0]]])

mapped_fn(indices)
```
Thanks for the great reproducer! Thanks to that, this just took 2 mins to spot once I ran your repro.

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -925,7 +925,7 @@ def replicate(val, axis_size, nrep, devices=None, backend=None):
     A ShardedDeviceArray of length `axis_size` where each shard is equal to
     ``val``.
   """
-  device_count = (len(devices) if devices else xb.local_device_count())
+  device_count = (len(devices) if devices else xb.local_device_count(backend))
   if nrep > device_count:
     msg = ("Cannot replicate across %d replicas because only %d local devices "
            "are available." % (nrep, device_count))


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1481,6 +1481,15 @@ def pmapped_multi_step(state):
     u = np.ones((device_count, 100))
     multi_step_pmap(u)  # doesn't crash
 
+  @jtu.skip_on_devices("cpu")
+  def test_replicate_backend(self):
+    # https://github.com/google/jax/issues/4223
+    def fn(indices):
+      return jnp.equal(indices, jnp.arange(3)).astype(jnp.float32)
+    mapped_fn = jax.pmap(fn, axis_name='i', backend='cpu')
+    mapped_fn = jax.pmap(mapped_fn, axis_name='j', backend='cpu')
+    indices = np.array([[[2], [1]], [[0], [0]]])
+    mapped_fn(indices)  # doesn't crash
 
 
 class VmapOfPmapTest(jtu.JaxTestCase):
Repository: google/jax

Issue:
Ppermute batching errors
This code.

```python
import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.pmap, axis_name="i")
@jax.vmap
def f(a):
  jax.lax.ppermute(a, "i", [0, 1, 2, 3, 4, 5, 6, 7, 8])

# Run it.
f(jnp.arange(8 * 8 * 8).reshape((8, 8, 8)))
```

Dies with this
```
TypeError: _allreduce_translation_rule() got an unexpected keyword argument 'perm'
```

I think this is because of this line
https://github.com/google/jax/blob/40e20242db0ed1d8cceb1d93b78d07c248e693a6/jax/lax/lax_parallel.py#L550

For some reason, the batching primative was set as `pmin`. 

I'll go ahead and make a PR to fix this.



Discussion:

Resolving patch:
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -547,7 +547,6 @@ def _ppermute_batcher(vals_in, dims_in, axis_size, axis_name, perm):
 ad.deflinear(ppermute_p, _ppermute_transpose_rule)
 xla.parallel_translations[ppermute_p] = _ppermute_translation_rule
 pxla.multi_host_supported_collectives.add(ppermute_p)
-batching.primitive_batchers[ppermute_p] = partial(_collective_batcher, pmin_p)
 batching.collective_rules[ppermute_p] = _ppermute_batcher
 
 


Test patch:
Repository: google/jax

Issue:
np.mean raises warning when input is float64 in jax 0.2.1
With the new update, the following script
```python
import numpy as np
import jax.numpy as jnp
jnp.mean(np.ones(30))  # or jnp.mean(1)
```
raises the warning
```
UserWarning: Explicitly requested dtype float64 requested in sum is not available,
and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64
configuration option or the JAX_ENABLE_X64 shell environment variable. See
https://github.com/google/jax#current-gotchas for more.
  warnings.warn(msg.format(dtype, fun_name , truncated_dtype))
```
I think this new behavior is unexpected so I would like to open this issue.


Discussion:
Thanks for the report. I believe the reason this started appearing in 0.2.1 is because of #4444. The best fix would be to call `dtypes.canonicalize_dtype()` here before passing the dtype to `sum()`: https://github.com/google/jax/blob/d4da9cc12cc785b0d174bdadddf29a5167239724/jax/numpy/lax_numpy.py#L1864-L1867

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1865,6 +1865,7 @@ def mean(a, axis=None, dtype=None, out=None, keepdims=False):
       dtype = float_
     else:
       dtype = _dtype(a)
+  dtype = dtypes.canonicalize_dtype(dtype)
 
   return lax.div(
       sum(a, axis, dtype=dtype, keepdims=keepdims),


Test patch: