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

Issue:
jvp docs names output v which is confuses with v in jvp
[The docs for `jvp`](https://jax.readthedocs.io/en/latest/jax.html?highlight=jvp#jax.jvp) currently have an example where the output is called `y, v`:
```
>>> y, v = jax.jvp(jax.numpy.sin, (0.1,), (0.2,))
>>> print(y)
0.09983342
>>> print(v)
0.19900084
```

this is confusing as it collides with the `v` in `jvp`. The output should be renamed to something useful. What's wrong with `primals, tangents_out` ? 


Discussion:

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2279,10 +2279,10 @@ def jvp(
 
   >>> import jax
   >>>
-  >>> y, v = jax.jvp(jax.numpy.sin, (0.1,), (0.2,))
-  >>> print(y)
+  >>> primals, tangents = jax.jvp(jax.numpy.sin, (0.1,), (0.2,))
+  >>> print(primals)
   0.09983342
-  >>> print(v)
+  >>> print(tangents)
   0.19900084
   """
   _check_callable(fun)


Test patch:
Repository: google/jax

Issue:
Broken links after moved design_notes folder
Looks like the `design_notes` folder was moved from the top level into `docs/` some time ago, but there are still a handful of broken links around from this change. 

Example of one:
https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#JAX-PRNG
The link to the design of JAX's prng is broken. 



Discussion:

Resolving patch:
diff --git a/jax/_src/config.py b/jax/_src/config.py
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -175,18 +175,18 @@ def parse_flags_with_absl(self):
       if not FLAGS.jax_omnistaging:
         raise Exception(
           "Disabling of omnistaging is no longer supported in JAX version 0.2.12 and higher: "
-          "see https://github.com/google/jax/blob/main/design_notes/omnistaging.md.\n"
+          "see https://github.com/google/jax/blob/main/docs/design_notes/omnistaging.md.\n"
           "To remove this warning, unset the JAX_OMNISTAGING environment variable.")
 
   def enable_omnistaging(self):
     warnings.warn(
         "enable_omnistaging() is a no-op in JAX versions 0.2.12 and higher;\n"
-        "see https://github.com/google/jax/blob/main/design_notes/omnistaging.md")
+        "see https://github.com/google/jax/blob/main/docs/design_notes/omnistaging.md")
 
   def disable_omnistaging(self):
     raise Exception(
       "Disabling of omnistaging is no longer supported in JAX version 0.2.12 and higher: "
-      "see https://github.com/google/jax/blob/main/design_notes/omnistaging.md.")
+      "see https://github.com/google/jax/blob/main/docs/design_notes/omnistaging.md.")
 
   def define_bool_state(
     self, name: str, default: bool, help: str, *,


Test patch:
Repository: google/jax

Issue:
Implement `.flat` on JAX arrays
JAX arrays are missing the `.flat` property:
https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flat.html

We should add it!


Discussion:
Just a note on this: the full `np.flatiter` API is more than just an iterator, it also supports advanced indexing operations. For example:
```python
In [1]: import numpy as np                                                                                                     

In [2]: x = np.arange(5)                                                                                                       

In [3]: mask = x > 2                                                                                                           

In [4]: x.flat[mask]                                                                                                           
Out[4]: array([3, 4])

In [5]: x.flat[[1, 2, 3]]                                                                                                     
Out[5]: array([1, 2, 3])
```
I think the implementation of this is going to have some subtleties, particularly with regard to JIT lowering and batching.

The easiest way to tick all the boxes might be to do `flat = property(ravel)`
+1 for `flat = property(ravel)`. I don't think we want to make any attempt at implementing NumPy's iterator API.

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
@@ -4553,6 +4553,10 @@ def _view(arr, dtype=None, type=None):
     return lax.bitcast_convert_type(arr_bytes, uint8).astype(dtype)
   return lax.bitcast_convert_type(arr_bytes, dtype)
 
+def _notimplemented_flat(self):
+  raise NotImplementedError("JAX DeviceArrays do not implement the arr.flat property: "
+                            "consider arr.flatten() instead.")
+
 ### track unimplemented functions
 
 _NOT_IMPLEMENTED_DESC = """
@@ -5010,6 +5014,7 @@ def _set_shaped_array_attributes(shaped_array):
   setattr(shaped_array, "reshape", core.aval_method(_reshape))
   setattr(shaped_array, "transpose", core.aval_method(_transpose))
   setattr(shaped_array, "flatten", core.aval_method(ravel))
+  setattr(shaped_array, "flat", core.aval_property(_notimplemented_flat))
   setattr(shaped_array, "T", core.aval_property(transpose))
   setattr(shaped_array, "real", core.aval_property(real))
   setattr(shaped_array, "imag", core.aval_property(imag))
@@ -5043,6 +5048,7 @@ def _set_device_array_base_attributes(device_array):
   setattr(device_array, "reshape", _reshape)
   setattr(device_array, "transpose", _transpose)
   setattr(device_array, "flatten", ravel)
+  setattr(device_array, "flat", property(_notimplemented_flat))
   setattr(device_array, "T", property(transpose))
   setattr(device_array, "real", property(real))
   setattr(device_array, "imag", property(imag))


Test patch:
Repository: google/jax

Issue:
hit assertion error in api._get_axis_size
A user hit [this assertion](https://github.com/google/jax/blob/f6f6b4f04cb1cfb64d531eb41172e1783c57fe47/jax/_src/api.py#L1562). We haven't been able to make a minimal repro yet.


Discussion:

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1593,10 +1593,15 @@ def _get_axis_size(name: str, shape: Tuple[int, ...], axis: int):
     # in which case we can produce an error message based on argument indices,
     # or if it has nested containers.
     if kws:
-      # if keyword arguments are included in the tree, we make adapt the error
+      position_only_tree, leaf = treedef_children(tree)
+      if not treedef_is_leaf(leaf):
+        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(f"the tree of axis sizes is:\n{sizes}")) from None
+      # if keyword arguments are included in the tree, we adapt the error
       # message only to be about the positional arguments
-      tree, leaf = treedef_children(tree)
-      assert treedef_is_leaf(leaf)
+      tree = position_only_tree
+
     # TODO(mattjj,phawkins): add a way to inspect pytree kind more directly
     if tree == tree_flatten((0,) * tree.num_leaves)[1]:
       lines1 = [f"arg {i} has shape {np.shape(x)} and axis {d} is to be mapped"


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2586,6 +2586,18 @@ def test_vmap_unbatched_object_passthrough_issue_183(self):
     ans = vfun(lambda x: x + 1, jnp.arange(3))
     self.assertAllClose(ans, np.arange(1, 4), check_dtypes=False)
 
+  def test_vmap_mismatched_keyword(self):
+    # https://github.com/google/jax/issues/10193
+    @jax.vmap
+    def f(x, y):
+      return x + y
+
+    with self.assertRaisesRegex(
+        ValueError, "vmap got inconsistent sizes for array axes to be mapped:\n"
+        "the tree of axis sizes is:\n"
+        r"\(\(1,\), \{'y': 2\}\)"):
+      f(jnp.array([1]), y=jnp.array([1, 2]))
+
   def test_vmap_mismatched_axis_sizes_error_message_issue_705(self):
     # https://github.com/google/jax/issues/705
     def h(a, b):
Repository: google/jax

Issue:
improve error message for jnp.diagonal
```python
import jax.numpy as jnp
import numpy as np
M = jnp.ones(5)
```

Then, `np.diagonal(M)` returns a clear error message `ValueError: diag requires an array of at least two dimensions` whereas `jnp.diagonal(M)` returns a still-informative-but-less-clear `ValueError: axis 1 is out of bounds for array of dimension 1`.

Could we possibly change this behavior?


Discussion:
Thanks for the report - we should be able to improve that 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
@@ -2516,6 +2516,8 @@ def diag_indices_from(arr):
 def diagonal(a, offset=0, axis1: int = 0, axis2: int = 1):
   _check_arraylike("diagonal", a)
   a_shape = shape(a)
+  if ndim(a) < 2:
+    raise ValueError("diagonal requires an array of at least two dimensions.")
   offset = core.concrete_or_error(operator.index, offset, "'offset' argument of jnp.diagonal()")
 
   a = moveaxis(a, (axis1, axis2), (-2, -1))


Test patch:
Repository: google/jax

Issue:
`lax.rem` has wrong gradient for `y` when `x<0`
`lax.rem` has wrong gradient when `x<0`. For example,

```py
import jax

def fn(y):
    x = jax.numpy.array(-19., jax.numpy.float64)
    return jax.lax.rem(x, y)

y = jax.numpy.array(3., jax.numpy.float64)
res = fn(y)
print(res) # -1.0

a = jax.jacfwd(fn)(y)
print(a) # 7.0
```

Actually, the gradient should be 6 instead of 7. Because `-19 = (-6)*3 + (-1)`. When the `y` is `3.1`, the results will be `-0.4`. Thus, the gradient should be 6.

By the way, another two similar APIs `numpy.remainder` and `numpy.mod` also have this issue.

```py
def fn(y):
    x = jax.numpy.array(-19., jax.numpy.float64)
    return jax.numpy.remainder(x, y)

y = jax.numpy.array(3., jax.numpy.float64)
res = fn(y)
print(res) # 2.0

a = jax.jacfwd(fn)(y)
print(a) # 8.0
```

In these two APIs, `-19 = (-7)*3 + 2`, so their gradient should be 7 instead of 8.


Discussion:
Thanks for the report – I have the fix in #11469

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
@@ -2161,7 +2161,7 @@ def _div_transpose_rule(cotangent, x, y):
 ad.defjvp(
     rem_p,
     lambda g, x, y: _maybe_broadcast(broadcast_shapes(np.shape(x), np.shape(y)), g),
-    lambda g, x, y: mul(neg(g), floor(div(x, y))))
+    lambda g, x, y: mul(neg(g), mul(sign(div(x, y)), floor(abs(div(x, y))))))
 mlir.register_lowering(rem_p, partial(_nary_lower_mhlo, mhlo.RemOp))
 
 def _minmax_complex_lowering(x, y, *, lax_cmp_pick_x):


Test patch:
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -1060,19 +1060,24 @@ def f2(x, y):
 
   # TODO(mattjj): make this a more systematic test
   def testRemainder(self):
+    def gen_x(rng, size):
+      return rng.uniform(-9, 9, size=size)
+
+    def gen_y(rng, size):
+      # avoid values near zero because gradients diverge
+      return rng.uniform(0.1, 5, size=size) * rng.choice([-1, 1], size=size)
+
     rng = self.rng()
-    x = rng.uniform(-0.9, 9, size=(3, 4))
-    y = rng.uniform(0.7, 1.9, size=(3, 1))
+    x = gen_x(rng, (5, 8))
+    y = gen_y(rng, (1, 8))
     assert not set(np.unique(x)) & set(np.unique(y))
-    # TODO(jakevdp) try to make these tolerances tighter.
-    tol = 1e-1
-    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"], tol, tol)
+    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"])
 
     rng = self.rng()
-    x = rng.uniform(-0.9, 9, size=(1, 4))
-    y = rng.uniform(0.7, 1.9, size=(3, 4))
+    x = gen_x(rng, (1, 8))
+    y = gen_y(rng, (5, 8))
     assert not set(np.unique(x)) & set(np.unique(y))
-    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"], tol, tol)
+    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"])
 
   def testHigherOrderGradientOfReciprocal(self):
     # Regression test for https://github.com/google/jax/issues/3136
Repository: google/jax

Issue:
tree_util deprecation warnings should set warning class and stacklevel
If I try using deprecated functions like `jax.tree_map` using the development version of JAX, I see messages like:
```
>>> jax.tree_map(lambda x: x, [1, 2, 3])
/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py:529: UserWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
  warnings.warn(f"jax.{f.__name__} is deprecated, and will be removed in a future release. "
[1, 2, 3]
```

This is OK, but a better warning would set `category` (so it can be filtered appropriately) and `stacklevel` (so it's clear to users where it's coming from).

Either `DeprecationWarning` or `FutureWarning` would be appropriate for the category, depending on how "noisy" you want the warnings to be. I would probably lean towards `FutureWarning` for JAX, since JAX is typically used by developers rather than end users.


Discussion:
Thanks - that was an oversight on my part. Fixed in #11507

Resolving patch:
diff --git a/jax/_src/tree_util.py b/jax/_src/tree_util.py
--- a/jax/_src/tree_util.py
+++ b/jax/_src/tree_util.py
@@ -527,7 +527,8 @@ def _deprecate(f):
   @functools.wraps(f)
   def wrapped(*args, **kwargs):
     warnings.warn(f"jax.{f.__name__} is deprecated, and will be removed in a future release. "
-                  f"Use jax.tree_util.{f.__name__} instead.")
+                  f"Use jax.tree_util.{f.__name__} instead.",
+                  category=FutureWarning, stacklevel=2)
     return f(*args, **kwargs)
   return wrapped
 


Test patch:
Repository: google/jax

Issue:
BUG: jnp.linalg.matrix_rank does not operate on stacks of matrices
### Description

```
myInput = np.random.normal(size=(9,10,10))
myInput.shape
# (9,10,10)
np.linalg.matrix_rank(myInput)
# array([10, 10, 10, 10, 10, 10, 10, 10, 10])
```

jax fails on this same stack:
```
jax.numpy.linalg.matrix_rank(jnp.array(myInput))
### TypeError: array should have 2 or fewer dimensions
### Expected: DeviceArray([10, 10, 10, 10, 10, 10, 10, 10, 10])
```

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

jax 0.3.13, jaxlib 0.3.10

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

- [ ] CPU
- [X] GPU
- [ ] TPU

### Additional System Info

_No response_


Discussion:
i haven't tested, but suspect this might apply to the jax version of `numpy.linalg.svd` as well.
Thanks for the report - I'll take a look.

Resolving patch:
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -111,14 +111,13 @@ def matrix_power(a, n):
 @jit
 def matrix_rank(M, tol=None):
   M, = _promote_dtypes_inexact(jnp.asarray(M))
-  if M.ndim > 2:
-    raise TypeError("array should have 2 or fewer dimensions")
   if M.ndim < 2:
     return jnp.any(M != 0).astype(jnp.int32)
   S = svd(M, full_matrices=False, compute_uv=False)
   if tol is None:
-    tol = S.max() * np.max(M.shape).astype(S.dtype) * jnp.finfo(S.dtype).eps
-  return jnp.sum(S > tol)
+    tol = S.max(-1) * np.max(M.shape[-2:]).astype(S.dtype) * jnp.finfo(S.dtype).eps
+  tol = jnp.expand_dims(tol, np.ndim(tol))
+  return jnp.sum(S > tol, axis=-1)
 
 
 @custom_jvp


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -924,7 +924,8 @@ def testMatrixPower(self, shape, dtype, n):
       {"testcase_name": "_shape={}".format(
            jtu.format_shape_dtype_string(shape, dtype)),
        "shape": shape, "dtype": dtype}
-      for shape in [(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50)]
+      for shape in [(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50),
+                    (3, 4, 5), (2, 3, 4, 5)]
       for dtype in float_types + complex_types))
   def testMatrixRank(self, shape, dtype):
     rng = jtu.rand_default(self.rng())
Repository: google/jax

Issue:
BUG: numpy and jax differ on modulus by zero
### Description

Jax and numpy differ in their behavior with regards to computing a modulus by zero. Jax treats it as a no-op, and numpy treats it as the zero function.

## Numpy

```python
Python 3.10.5 (main, Jul 15 2022, 03:56:49) [GCC 11.3.0] on linux
>>> import numpy as np
>>> a = np.array([1,2,3], dtype=np.uint32)
>>> np.mod(a, 0)
<stdin>:1: RuntimeWarning: divide by zero encountered in remainder
array([0, 0, 0], dtype=uint32)
```

## JAX

```python
>>> import jax.numpy as jnp
>>> a = jnp.array([1,2,3], dtype=jnp.uint32)
>>> jnp.mod(a, 0)
DeviceArray([1, 2, 3], dtype=uint32)
```

### What jax/jaxlib version are you using?

HEAD

### Which accelerator(s) are you using?

CPU

### Additional System Info

Python 3.9, 3.10.5


Discussion:
Thanks for filing - I noticed also that this leads to strange behaviors in `divmod`:
```python
In [2]: print(*np.divmod(2, 0))                                                                                                                                                  
<ipython-input-28-e495b2cbe1c8>:1: RuntimeWarning: divide by zero encountered in divmod
0 0

In [3]: print(*jnp.divmod(2, 0))                                                                                                                                                 
-2 2
```

Resolving patch:
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -529,6 +529,8 @@ def frexp(x):
 def remainder(x1, x2):
   x1, x2 = _promote_args_numeric("remainder", x1, x2)
   zero = _constant_like(x1, 0)
+  if dtypes.issubdtype(x2.dtype, np.integer):
+    x2 = _where(x2 == 0, lax_internal._ones(x2), x2)
   trunc_mod = lax.rem(x1, x2)
   trunc_mod_not_zero = lax.ne(trunc_mod, zero)
   do_plus = lax.bitwise_and(


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
@@ -286,9 +286,9 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     op_record("rad2deg", 1, float_dtypes, all_shapes, jtu.rand_default, []),
     op_record("ravel", 1, all_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("real", 1, number_dtypes, all_shapes, jtu.rand_some_inf, []),
-    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero, [],
+    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_some_zero, [],
               tolerance={np.float16: 1e-2}),
-    op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_nonzero, []),
+    op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_some_zero, []),
     op_record("modf", 1, float_dtypes, all_shapes, jtu.rand_default, []),
     op_record("modf", 1, int_dtypes + unsigned_dtypes, all_shapes,
               jtu.rand_default, [], check_dtypes=False),
Repository: google/jax

Issue:
pure_callback passes jax.DeviceArray to the callback on CPU when not jitted
### Description

As title. 
This only happens when the callback is executed outside of jit boundaries.

This is different from the documented behaviour.

```python
➜ python
Python 3.10.6 (main, Aug 23 2022, 11:35:18) [Clang 13.1.6 (clang-1316.0.21.2.5)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> def test(x):
...   print(type(x))
...   return x
...
>>> def f(x):
...   return jax.pure_callback(test, x, x)
...
>>> x= jax.numpy.ones(3)
>>> f(x)
<class 'jaxlib.xla_extension.DeviceArray'>
DeviceArray([1., 1., 1.], dtype=float32)
>>> jax.jit(f)(x)
<class 'numpy.ndarray'>
DeviceArray([1., 1., 1.], dtype=float32)
```

### What jax/jaxlib version are you using?

jax 0.3.17

### Which accelerator(s) are you using?

CPU

### Additional System Info

MacOs


Discussion:
This is a straightforward fix on our end, though I'm curious, does this cause problems for you given that `DeviceArray`s behave like `np.ndarray`s?

Resolving patch:
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -14,6 +14,8 @@
 """Module for JAX callbacks."""
 from __future__ import annotations
 
+import functools
+
 from typing import Any, Callable, Sequence
 
 from jax import core
@@ -21,6 +23,7 @@
 from jax._src import dtypes
 from jax._src import lib as jaxlib
 from jax._src import util
+from jax._src import dispatch
 from jax.interpreters import ad
 from jax.interpreters import batching
 from jax.interpreters import mlir
@@ -33,11 +36,12 @@
 map, unsafe_map = util.safe_map, map
 
 
-@pure_callback_p.def_impl
 def pure_callback_impl(*args, result_avals, callback: Callable[..., Any],
                        vectorized: bool):
   del vectorized, result_avals
   return callback(*args)
+pure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,
+                                           pure_callback_p))
 
 
 @pure_callback_p.def_abstract_eval
@@ -102,7 +106,7 @@ def pure_callback_lowering(ctx, *args, callback, **params):
                               "Please upgrade to a jaxlib >= 0.3.15.")
 
   def _callback(*flat_args):
-    return tuple(pure_callback_p.impl(*flat_args, callback=callback, **params))
+    return tuple(pure_callback_impl(*flat_args, callback=callback, **params))
 
   result, _, keepalive = mlir.emit_python_callback(
       ctx, _callback, None, list(args), ctx.avals_in, ctx.avals_out, False,


Test patch:
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -469,6 +469,17 @@ def tearDown(self):
     super().tearDown()
     dispatch.runtime_tokens.clear()
 
+  @jtu.skip_on_devices(*disabled_backends)
+  def test_pure_callback_passes_ndarrays_without_jit(self):
+
+    def cb(x):
+      self.assertIs(type(x), np.ndarray)
+      return x
+
+    def f(x):
+      return jax.pure_callback(cb, x, x)
+    f(jnp.array(2.))
+
   @jtu.skip_on_devices(*disabled_backends)
   def test_simple_pure_callback(self):
Repository: google/jax

Issue:
Misleading error message in broadcast
### Description

```python
In [64]: x, y = jnp.ones((2, 4)), jnp.ones((2,))
In [65]: x + y
....
/jax/_src/numpy/util.py in _promote_shapes(fun_name, *args)
    246         if config.jax_numpy_rank_promotion != "allow":
    247           _rank_promotion_warning_or_error(fun_name, shapes)
--> 248         result_rank = len(lax.broadcast_shapes(*shapes))
    249         return [_broadcast_to(arg, (1,) * (result_rank - len(shp)) + shp)
    250                 for arg, shp in zip(args, shapes)]

ValueError: Incompatible shapes for broadcasting: ((2, 4), (1, 2))
```
```python
In [66]: xx, yy = np.ones((2, 4)), np.ones((2,))
In [67]: xx + yy
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-67-9e546835689c> in <module>()
----> 1 xx + yy

ValueError: operands could not be broadcast together with shapes (2,4) (2,)
```

The error message from `jnp` does not tell the shape given by user and that's misleading for debugging. `np` has a better behavior here.

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

latest

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

_No response_

### Additional System Info

_No response_


Discussion:
Thanks for the report - we should be able to improve that

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
@@ -149,8 +149,7 @@ def _broadcast_shapes_uncached(*shapes):
   shape_list = [(1,) * (ndim - len(shape)) + shape for shape in shapes]
   result_shape = _try_broadcast_shapes(shape_list)
   if result_shape is None:
-    raise ValueError("Incompatible shapes for broadcasting: {}"
-                     .format(tuple(shape_list)))
+    raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
   return result_shape
 
 def _broadcast_ranks(s1, s2):


Test patch: