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

Issue:
Consider setting optimize=True as default for einsum
We should consider changing the default of einsum to be optimized, as the search for the best contraction order will generally be faster than compilation anyway.

The contraction order calculation tends to only be slow when there are >10 tensors, which doesn't happen very often.


Discussion:
Nice idea! What's numpy do here?
Numpy defaults to False

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
@@ -2558,7 +2558,7 @@ def tensordot(a, b, axes=2, precision=None):
 
 @_wraps(onp.einsum, lax_description=_PRECISION_DOC)
 def einsum(*operands, **kwargs):
-  optimize = kwargs.pop('optimize', 'auto')
+  optimize = kwargs.pop('optimize', True)
   optimize = 'greedy' if optimize is True else optimize
   precision = kwargs.pop('precision', None)
   if kwargs:


Test patch:
Repository: google/jax

Issue:
Expose stable flag on XLA sort
XLA's sort HLOs have an "is_stable" attribute, which defaults to false; in practice, the sort used on CPU is stable but the one on GPU isn't. This means that e.g. this code will fail on GPU:
```python
x = onp.array([5, 1, 2, 6, 5, -2, 2, 0, 0, 1, 0, 4])
assert jnp.argsort(x) == onp.argsort(x)
```
and #2760 won't match NumPy behavior (ties to earliest) either.


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -4494,6 +4494,7 @@ def _sort_batch_rule(batched_args, batch_dims, *, dimension):
 
 sort_p = standard_primitive(sort_shape, _input_dtype, 'sort')
 ad.defjvp(sort_p, _sort_jvp_rule)
+xla.translations[sort_p] = partial(standard_translate, 'sort', is_stable=True)
 batching.primitive_batchers[sort_p] = _sort_batch_rule
 
 def _sort_key_val_abstract_eval(keys, values, *, dimension):
@@ -4562,7 +4563,8 @@ def _sort_key_val_batch_rule(batched_args, batch_dims, *, dimension):
 sort_key_val_p.multiple_results = True
 sort_key_val_p.def_impl(partial(xla.apply_primitive, sort_key_val_p))
 sort_key_val_p.def_abstract_eval(_sort_key_val_abstract_eval)
-xla.translations[sort_key_val_p] = partial(standard_translate, 'sort_key_val')
+xla.translations[sort_key_val_p] = partial(standard_translate, 'sort_key_val',
+                                           is_stable=True)
 ad.primitive_jvps[sort_key_val_p] = _sort_key_val_jvp
 ad.primitive_transposes[sort_key_val_p] = _sort_key_val_transpose_rule
 batching.primitive_batchers[sort_key_val_p] = _sort_key_val_batch_rule


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
@@ -1095,7 +1095,6 @@ def testRepeat(self, axis, shape, dtype, repeats, rng_factory):
       for return_index in [False, True]
       for return_inverse in [False, True]
       for return_counts in [False, True]))
-  @jtu.skip_on_devices("gpu", "tpu")  # https://github.com/google/jax/issues/2779
   def testUnique(self, shape, dtype, return_index, return_inverse, return_counts, rng):
     args_maker = lambda: [rng(shape, dtype)]
     onp_fun = lambda x: onp.unique(x, return_index, return_inverse, return_counts)
diff --git a/tests/lax_scipy_sparse_test.py b/tests/lax_scipy_sparse_test.py
--- a/tests/lax_scipy_sparse_test.py
+++ b/tests/lax_scipy_sparse_test.py
@@ -101,12 +101,14 @@ def args_maker():
         check_dtypes=True,
         tol=3e-5)
 
+    # TODO(shoyer,mattjj): I had to loosen the tolerance for complex64[7,7]
+    # with preconditioner=random
     self._CheckAgainstNumpy(
         partial(scipy_cg, M=M, maxiter=3),
         partial(lax_cg, M=M, maxiter=3),
         args_maker,
         check_dtypes=True,
-        tol=1e-4)
+        tol=3e-3)
 
     self._CheckAgainstNumpy(
         np.linalg.solve,
Repository: google/jax

Issue:
ShardedDeviceArray access should be threadsafe
This used to be the case, but I broke it in https://github.com/google/jax/commit/07571ae4dd3fceee580aa49c4490f99ce7f6b6de here: https://github.com/google/jax/blob/master/jax/interpreters/pxla.py#L499 (self._npy_value may be accessed by another thread before it's populated). I'm planning to fix this today or Monday.


Discussion:

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -496,10 +496,11 @@ def block_until_ready(self):
   def _value(self):
     if self._npy_value is None:
       self.copy_to_host_async()
-      self._npy_value = onp.empty(self.aval.shape, self.aval.dtype)
+      npy_value = onp.empty(self.aval.shape, self.aval.dtype)
       for i in range(0, len(self.device_buffers),
                      self.sharding_spec.replication_factor):
-        self._npy_value[self.indices[i]] = self.device_buffers[i].to_py()
+        npy_value[self.indices[i]] = self.device_buffers[i].to_py()
+      self._npy_value = npy_value
     return self._npy_value
 
   def __getitem__(self, idx):


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -13,9 +13,11 @@
 # limitations under the License.
 
 
+from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 import os
 from random import shuffle
+import threading
 from unittest import SkipTest
 
 import numpy as onp
@@ -1147,6 +1149,35 @@ def f(x, y):
     self.assertAllClose(ans, expected, check_dtypes=False)
 
 
+class ShardedDeviceArrayTest(jtu.JaxTestCase):
+
+  def testThreadsafeIndexing(self):
+    # NOTE(skye): I picked these values to be big enough to cause interesting
+    # execution overlap, but small enough to not use too much memory. YMMV.
+    shape = (8, 8000, 1000)
+
+    if jax.device_count() < shape[0]:
+      raise SkipTest(f"requires {shape[0]} devices")
+
+    x = np.arange(np.prod(shape)).reshape(shape)
+    sharded_x = pmap(lambda x: x)(x)
+
+    num_threads = 10
+    futures = []
+    expected = []
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+      for i in range(num_threads):
+        idx = i % shape[0]
+        # Mix together different kinds of indices
+        if i % 2 == 0:
+          idx = slice(idx, idx + 1)
+        futures.append(executor.submit(
+            lambda: [sharded_x[idx] for _ in range(10)][0]))
+        expected.append(x[idx])
+      actual = [f.result() for f in futures]
+    self.assertAllClose(actual, expected, check_dtypes=False)
+
+
 class SpecToIndicesTest(jtu.JaxTestCase):
 
   def testShardsPerAxis(self):
Repository: google/jax

Issue:
np.array named columns and different column types not working
Classical numpy arrays allow to specify column names/identifiers and assign different types to columns.
From the documentation:

> Data-type consisting of more than one element:
>
> ```
>>>> x = np.array([(1,2),(3,4)],dtype=[('a','<i4'),('b','<i4')])
>>>> x['a']
>array([1, 3])
>  ```

The same documentation is given for jax's `np.array` but trying the above example yields a
```
TypeError: unhashable type: 'list'
```
for the `dtype` argument (raised in function `_check_user_dtype_supported`,  line 5096, `jax/lax/lax.py`).
Is there a chance that this kind of array specifications could be supported? If not, could at least the documentation be clarified?


Discussion:
JAX doesn’t support structured dtypes, but I agree that we should provide better a error message here.
too bad, but thanks for the clarification!

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -5173,8 +5173,13 @@ def _abstractify(x):
   return raise_to_shaped(core.get_aval(x))
 
 
+
 def _check_user_dtype_supported(dtype, fun_name=None):
-  if dtype is not None and onp.dtype(dtype) != dtypes.canonicalize_dtype(dtype):
+  onp_dtype = onp.dtype(dtype)
+  if onp_dtype.kind not in "biufc" and onp_dtype.type != dtypes.bfloat16:
+    msg = f"JAX only supports number and bool dtypes, got dtype {dtype}"
+    raise TypeError(msg)
+  if dtype is not None and onp_dtype != dtypes.canonicalize_dtype(dtype):
     msg = ("Explicitly requested dtype {} {} is not available, "
            "and will be truncated to dtype {}. To enable more dtypes, set the "
            "jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell "


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
@@ -1728,6 +1728,11 @@ def testArray(self, arg, ndmin, dtype):
     self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
 
+  def testArrayUnsupportedDtypeError(self):
+    with self.assertRaisesRegex(TypeError,
+                                "JAX only supports number and bool dtypes.*"):
+      jnp.array(3, [('a','<i4'),('b','<i4')])
+
   def testIssue121(self):
     assert not onp.isscalar(jnp.array(3))
Repository: google/jax

Issue:
Argmax on empty array returns 2147483647
How to reproduce:

```
>>> jnp.argmax(jnp.array([]))
DeviceArray(2147483647, dtype=int32)
```

Confirmed to occur on CPU and TPU.

For reference, in numpy an exception is returned: "ValueError: attempt to get argmax of an empty sequence".


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
@@ -2827,7 +2827,7 @@ def argmax(a, axis=None):
   if axis is None:
     a = ravel(a)
     axis = 0
-  return _argminmax(max, a, axis)
+  return _argminmax("argmax", max, a, axis)
 
 
 _NANARG_DOC = """\
@@ -2850,7 +2850,7 @@ def argmin(a, axis=None):
   if axis is None:
     a = ravel(a)
     axis = 0
-  return _argminmax(min, a, axis)
+  return _argminmax("argmin", min, a, axis)
 
 
 @_wraps(onp.nanargmin, lax_description=_NANARG_DOC.format("min"))
@@ -2864,7 +2864,9 @@ def nanargmin(a, axis=None):
 
 
 # TODO(mattjj): redo this lowering with a call to variadic lax.reduce
-def _argminmax(op, a, axis):
+def _argminmax(name, op, a, axis):
+  if size(a) == 0:
+    raise ValueError("attempt to get {} of an empty sequence".format(name))
   shape = [1] * a.ndim
   shape[axis] = a.shape[axis]
   idxs = lax.tie_in(a, arange(a.shape[axis])).reshape(shape)


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
@@ -696,6 +696,16 @@ def jnp_fun(array_to_reduce):
         raise
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": rec.test_name.capitalize(),
+       "name": rec.name, "jnp_op": getattr(jnp, rec.name)}
+      for rec in JAX_ARGMINMAX_RECORDS))
+  def testArgMinMaxEmpty(self, name, jnp_op):
+    name = name[3:] if name.startswith("nan") else name
+    msg = "attempt to get {} of an empty sequence".format(name)
+    with self.assertRaises(ValueError, msg=msg):
+      jnp_op(onp.array([]))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_{}_{}".format(
           jtu.format_shape_dtype_string(lhs_shape, lhs_dtype),
Repository: google/jax

Issue:
random.logistic returns -inf
If I call jax.random.logistic with key array([2308802225, 3876068305], dtype=uint32) and shape=(10000, 12), the sampled array contains -inf. I don't think this behavior is correct or at least it is not expected and causes downstream code that assumes finite floats to crash.

I can work around this, but I think it's worth fixing. Thanks!


Discussion:
This is a consequence of using a 32-bit implementation of the standard logit generating function: x is drawn from a uniform distribution between 0 and 1, then transformed according to `y = log(x / (1 - x))`. For real numbers there is an infinitessimally small chance of `x=0` which maps to `y=-inf`, in 32-bit arithmetic the chance is approximately 1 in 2^30.

I'll think about how to best address this.
Wait, it's even worse than that, because ``uniform()`` only randomizes the mantissa: this means the chance is about a 1 in 2^23 of getting -inf; that's confirmed empirically:
```
>>> (uniform(PRNGKey(0), (100000000,)) == 0).sum()
DeviceArray(19, dtype=int32)
```
I think the best approach would be to modify the generating algorithm to use 32-bit unsigned integers rather than 32-bit uniform floats; then the chances of hitting this would be one in 4 billion, rather than 1 in 8 million as it is currently.
Could you avoid the problem by drawing from the range `(0, 1)` instead of `[0, 1)` ?
Yes, that would fix the -inf issue. But it still might be preferable to use an approach that allows the result to have a cardinality greater than 2^23

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -1227,8 +1227,21 @@ def logistic(key, shape=(), dtype=onp.float64):
 
 @partial(jit, static_argnums=(1, 2))
 def _logistic(key, shape, dtype):
+  # Mathematically, we can compute the distribution by generating uniformly-distributed
+  # numbers x in the open interval (a, b) and computing:
+  #   z = log[ (x - a) / (b - x))
+  # It's important to avoid x=a or x=b, which lead to infinite values for z.
+  # The uniform() function generates pseudorandom floating point numbers x in the
+  # semi-closed interval [0, 1), so if used directly  with (a,b)=(0,1), it will
+  # lead to infinite output in a small number of cases (as many as 1 in 2^23 for float32).
+  #
+  # Instead, we let (a, b) = (-ε, 1) where ε is the smallest step between floating point
+  # values: then numbers in the interval (-ε, 1) are approximated by standard uniformly
+  # drawn numbers in [0, 1).
   _check_shape("logistic", shape)
-  return logit(uniform(key, shape, dtype))
+  x = uniform(key, shape, dtype)
+  eps = np.finfo(dtype).eps
+  return lax.log(lax.div(lax.add(lax._const(x, eps), x), lax.sub(lax._const(x, 1), x)))
 
 
 def pareto(key, b, shape=None, dtype=onp.float64):


Test patch:
Repository: google/jax

Issue:
minor problem in error message for jax.experimental.optimizers.piecewise_constant
Hi,

I think there is a small problem with the error message for jax.experimental.optimizers.piecewise_constant.

```python
def piecewise_constant(boundaries, values):
  boundaries = jnp.array(boundaries)
  values = jnp.array(values)
  if not boundaries.ndim == values.ndim == 1:
    raise ValueError("boundaries and values must be sequences")
  if not boundaries.shape[0] == values.shape[0] - 1:
    raise ValueError("boundaries length must be one longer than values length")
```
I think the last line should read "boundaries length must be one less than values length" or something along those lines.

Cheers!


Discussion:

Resolving patch:
diff --git a/jax/experimental/optimizers.py b/jax/experimental/optimizers.py
--- a/jax/experimental/optimizers.py
+++ b/jax/experimental/optimizers.py
@@ -499,7 +499,7 @@ def piecewise_constant(boundaries, values):
   if not boundaries.ndim == values.ndim == 1:
     raise ValueError("boundaries and values must be sequences")
   if not boundaries.shape[0] == values.shape[0] - 1:
-    raise ValueError("boundaries length must be one longer than values length")
+    raise ValueError("boundaries length must be one shorter than values length")
 
   def schedule(i):
     return values[jnp.sum(i > boundaries)]


Test patch:
Repository: google/jax

Issue:
[jet] Unable to determine order from zeroseries
When `series_in` consists only of `zeroseries`, [this line](https://github.com/google/jax/blob/master/jax/experimental/jet.py#L116) will fail with `
ValueError: not enough values to unpack (expected 1, got 0)`.

This occurs when taking jet through `jax.nn.softplus = lambda x: jnp.logaddexp(x, 0)` as we pass one `zeroseries`to `convert_element_type`. We can avoid the error in this case by avoiding this primitive altogether and manually converting the dtype `lambda x: jnp.logaddexp(x, 0.)`, but I think it's better to fix this in general.


Discussion:

Resolving patch:
diff --git a/jax/experimental/jet.py b/jax/experimental/jet.py
--- a/jax/experimental/jet.py
+++ b/jax/experimental/jet.py
@@ -51,16 +51,17 @@ def flatten_fun_output(*args):
     yield tree_flatten(ans)
 
   f, out_tree = flatten_fun_output(lu.wrap_init(fun))
-  out_primals, out_terms = jet_fun(jet_subtrace(f)).call_wrapped(primals, series)
+  out_primals, out_terms = jet_fun(jet_subtrace(f), order).call_wrapped(primals, series)
   return tree_unflatten(out_tree(), out_primals), tree_unflatten(out_tree(), out_terms)
 
 @lu.transformation
-def jet_fun(primals, series):
+def jet_fun(order, primals, series):
   with core.new_master(JetTrace) as master:
+    master.order = order
     out_primals, out_terms = yield (master, primals, series), {}
     del master
-  out_terms = [tree_map(lambda x: onp.zeros_like(x, dtype=onp.result_type(out_primals[0])), series[0])
-               if s is zero_series else s for s in out_terms]
+  out_terms = [[onp.zeros_like(p)] * order if s is zero_series else s
+               for p, s in zip(out_primals, out_terms)]
   yield out_primals, out_terms
 
 @lu.transformation
@@ -112,8 +113,8 @@ def sublift(self, val):
 
   def process_primitive(self, primitive, tracers, params):
     assert not primitive.multiple_results  # TODO
+    order = self.master.order
     primals_in, series_in = unzip2((t.primal, t.terms) for t in tracers)
-    order, = {len(terms) for terms in series_in if terms is not zero_series}
     series_in = [[zero_term] * order if s is zero_series else s
                  for s in series_in]
     # TODO(mattjj): avoid always instantiating zeros


Test patch:
Repository: google/jax

Issue:
`issubdtype` errors when using dtype constructor
NumPy:
```python
np.issubdtype(np.int32, np.dtype('int32')) # ==> True
```
JAX NumPy:
```python
jnp.issubdtype(jnp.int32, jnp.dtype('int32')) # ==> TypeError: issubclass() arg 1 must be a class
```


Discussion:
Try:`jnp.issubdtype(jnp.int32, jnp.dtype('int32').type)`.  This is an unfortunate quirk of `numpy` whereby dtype objects are sometimes interchangeable with the types that they are based off.  See [here](https://github.com/numpy/numpy/issues/7242), for example.
Thanks for the pointer!
Is this issue solved?
No, I think it's still a bug. JAX should act like NumPy here.

Resolving patch:
diff --git a/jax/dtypes.py b/jax/dtypes.py
--- a/jax/dtypes.py
+++ b/jax/dtypes.py
@@ -125,12 +125,22 @@ def finfo(dtype):
   else:
     return np.finfo(dtype)
 
+def _issubclass(a, b):
+  """Determines if ``a`` is a subclass of ``b``.
+
+  Similar to issubclass, but returns False instead of an exception if `a` is not
+  a class.
+  """
+  try:
+    return issubclass(a, b)
+  except TypeError:
+    return False
 
 def issubdtype(a, b):
   if a == bfloat16:
     return b in [bfloat16, _bfloat16_dtype, np.floating, np.inexact,
                  np.number]
-  if not issubclass(b, np.generic):
+  if not _issubclass(b, np.generic):
     # Workaround for JAX scalar types. NumPy's issubdtype has a backward
     # compatibility behavior for the second argument of issubdtype that
     # interacts badly with JAX's custom scalar types. As a workaround,


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -149,6 +149,7 @@ def testIsSubdtype(self):
       self.assertTrue(dtypes.issubdtype(t, t))
       self.assertTrue(dtypes.issubdtype(np.dtype(t).type, t))
       self.assertTrue(dtypes.issubdtype(t, np.dtype(t).type))
+      self.assertTrue(dtypes.issubdtype(t, np.dtype(t)))
       if t != jnp.bfloat16:
         for category in [np.generic, jnp.inexact, jnp.integer, jnp.signedinteger,
                          jnp.unsignedinteger, jnp.floating, jnp.complexfloating]:
Repository: google/jax

Issue:
Request for better error message when pmap and psum don't match on axis_name.
```
a = np.arange(4)[np.newaxis,:]
print(a.shape)
def test(x):
  return jax.lax.psum(x, axis_name='batch')
jax.pmap(test)(a)
```
Gives an error message of `AttributeError: 'str' object has no attribute 'obj'`.
It took me forever to figure out that the issue was the pmap needed axis_name='batch'. It would be helpful if this gave a better error message here.


Discussion:
That might win an overall bad-JAX-error-message award! (Not a competition...)
But yeah, this is just a bug.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -1131,7 +1131,7 @@ def f_pmapped(*args, **kwargs):
   f_pmapped.__name__ = namestr(f_pmapped.__name__, axis_name)
   return f_pmapped
 
-class _TempAxisName(object):
+class _TempAxisName:
   def __init__(self, obj):
     self.obj = obj
   def __repr__(self):
@@ -1139,7 +1139,7 @@ def __repr__(self):
   def __hash__(self):
     return hash(self.obj)
   def __eq__(self, other):
-    return self.obj is other.obj
+    return type(other) is _TempAxisName and self.obj is other.obj
 
 
 def soft_pmap(fun: Callable, axis_name: Optional[AxisName] = None, *,
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -435,7 +435,7 @@ def check_backend_params(params, outer_backend):
   return {k: params[k] for k in params if k != 'backend'}
 
 
-class AxisEnv(object):
+class AxisEnv:
   def __init__(self, nreps, names=(), sizes=(), devices=None):
     assert isinstance(names, tuple)
     assert isinstance(sizes, tuple)
@@ -448,7 +448,10 @@ def extend_axis_env(env, name, size):
   return AxisEnv(env.nreps, env.names + (name,), env.sizes + (size,), env.devices)
 
 def axis_read(axis_env, axis_name):
-  return max(i for i, name in enumerate(axis_env.names) if name == axis_name)
+  try:
+    return max(i for i, name in enumerate(axis_env.names) if name == axis_name)
+  except ValueError:
+    raise NameError("unbound axis name: {}".format(axis_name))
 
 def axis_groups(axis_env, name):
   if isinstance(name, (list, tuple)):


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1184,6 +1184,15 @@ def matrix_vector(x, y, parallel=True):
     self.assertAllClose(result1, result3, check_dtypes=False, atol=1e-3, rtol=1e-3)
     self.assertAllClose(result1, result4, check_dtypes=False, atol=1e-3, rtol=1e-3)
 
+  def testPmapAxisNameError(self):
+    # https://github.com/google/jax/issues/3120
+    a = np.arange(4)[np.newaxis,:]
+    def test(x):
+      return jax.lax.psum(x, axis_name='batch')
+
+    with self.assertRaisesRegex(NameError, "unbound axis name: batch"):
+      jax.pmap(test)(a)
+
   def testPsumOnBooleanDtype(self):
     # https://github.com/google/jax/issues/3123
     n = xla_bridge.device_count()
Repository: google/jax

Issue:
broadcast_to docs are unhelpful
Right now the docs state "Like Numpy’s broadcast_to but doesn’t necessarily return views."
A few suggestions

1. Link to numpy's docs
2. Copy their docs in some way
3. Write our own version


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
@@ -1354,8 +1354,10 @@ def broadcast_arrays(*args):
   return [broadcast_to(arg, result_shape) for arg in args]
 
 
+@_wraps(np.broadcast_to, lax_description="""\
+The JAX version does not necessarily return a view of the input.
+""")
 def broadcast_to(arr, shape):
-  """Like Numpy's broadcast_to but doesn't necessarily return views."""
   arr = arr if isinstance(arr, ndarray) else array(arr)
   shape = canonicalize_shape(shape)  # check that shape is concrete
   arr_shape = _shape(arr)


Test patch: