nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: jax-ml/jax
Pull request #27449 (2025-03-25 22:47:31)

[sharding_in_types] Add out_sharding to `jax.random.normal`.
[sharding_in_types] Add out_sharding to `jax.random.normal`.

Drop into `Auto` mode inside for implementation.

Co-authored-by: Roy Frostig <frostig@google.com>



Patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
index 2d315ed0c..7277ed5aa 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -690,7 +690,8 @@ def choice(key: ArrayLike,
 
 def normal(key: ArrayLike,
            shape: Shape = (),
-           dtype: DTypeLikeFloat = float) -> Array:
+           dtype: DTypeLikeFloat = float,
+           out_sharding=None) -> Array:
   r"""Sample standard normal random values with given shape and float dtype.
 
   The values are returned according to the probability density function:
@@ -712,12 +713,16 @@ def normal(key: ArrayLike,
   """
   key, _ = _check_prng_key("normal", key)
   shape = core.canonicalize_shape(shape)
+  out_sharding = canonicalize_sharding(out_sharding, 'normal')
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.inexact):
     raise ValueError(f"dtype argument to `normal` must be a float or complex dtype, "
                      f"got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
-  return _normal(key, shape, dtype)
+  if out_sharding is None:
+    return _normal(key, shape, dtype)
+  return auto_axes(partial(_normal, shape=shape, dtype=dtype),
+                   out_shardings=out_sharding)(key)
 
 @partial(jit, static_argnums=(1, 2))
 def _normal(key, shape, dtype) -> Array:


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 528384358..d6673c6b6 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7245,6 +7245,25 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     else:
       self.assertIn('mhlo.sharding = "{devices=[2,2]<=[4]}"}', lowered_text)
 
+  @jtu.with_user_mesh((2, 2), ('x', 'y'))
+  def test_random_normal(self, mesh):
+    @jax.jit
+    def f(key):
+      out = jax.random.normal(key, shape=(8, 12), out_sharding=P('x', 'y'))
+      self.assertEqual(out.aval.sharding.spec, P('x', 'y'))
+      return out
+
+    key = jax.random.key(1)
+    out = f(key)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
+
+    lowered_text = f.lower(key).as_text()
+    if config.use_shardy_partitioner.value:
+      self.assertIn('sdy.sharding_constraint', lowered_text)
+      self.assertIn('<@mesh, [{"x"}, {"y"}]>', lowered_text)
+    else:
+      self.assertIn('mhlo.sharding = "{devices=[2,2]<=[4]}"}', lowered_text)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):
Repository: jax-ml/jax
Pull request #27491 (2025-03-26 20:17:55)

Add missing `__len__` to MutableArray
Add missing `__len__` to MutableArray

Fixes https://github.com/jax-ml/jax/issues/27476



Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index ef90341f5..14f3d9cc1 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1940,8 +1940,7 @@ class ShapedArray(UnshapedArray):
     if 'sharding' not in kwargs:
       kwargs['sharding'] = self.sharding
     if 'vma' not in kwargs:
-      kwargs['vma'] = getattr(self, 'vma',
-                                              frozenset())
+      kwargs['vma'] = getattr(self, 'vma', frozenset())
     return ShapedArray(shape, dtype, weak_type, **kwargs)
 
   ndim = property(lambda self: len(self.shape))
@@ -2170,6 +2169,7 @@ class MutableArray:
   def __getitem__(self, idx): return self._aval._getitem(self, idx)
   def __setitem__(self, idx, x): return self._aval._setitem(self, idx, x)
   def __repr__(self) -> str: return 'Mutable' + repr(self[...])
+  def __len__(self) -> int: return self._aval._len(self)
 pytype_aval_mappings[MutableArray] = lambda x: x._aval
 
 def mutable_array(init_val):
diff --git a/jax/_src/state/types.py b/jax/_src/state/types.py
index fa9d0cb9f..e926e3a35 100644
--- a/jax/_src/state/types.py
+++ b/jax/_src/state/types.py
@@ -333,6 +333,12 @@ class AbstractRef(core.AbstractValue):
   ndim = property(lambda self: len(self.shape))
   size = property(lambda self: math.prod(self.shape))
 
+  def _len(self, ignored_tracer) -> int:
+    try:
+      return self.shape[0]
+    except IndexError as err:
+      raise TypeError("len() of unsized object") from err  # same as numpy error
+
   @property
   def shape(self):
     try:


Test patch:
diff --git a/tests/mutable_array_test.py b/tests/mutable_array_test.py
index e962653ed..950bddf54 100644
--- a/tests/mutable_array_test.py
+++ b/tests/mutable_array_test.py
@@ -116,6 +116,18 @@ class MutableArrayTest(jtu.JaxTestCase):
                         check_dtypes=False)
     self.assertAllClose(w, 10, check_dtypes=False)
 
+  @parameterized.parameters([True, False])
+  def test_len_mutable_array(self, jit):
+    x_mut = core.mutable_array(jnp.zeros(3))
+
+    def f():
+      return jnp.int32(len(x_mut))
+
+    if jit:
+      f = jax.jit(f)
+
+    self.assertEqual(f(), 3)
+
   @parameterized.parameters([True, False])
   def test_internal_mutarray_basic(self, jit):
     def f():
Repository: jax-ml/jax
Pull request #27540 (2025-03-27 20:30:06)

jnp.power: support __jax_array__ on inputs


Patch:
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index e561b7ae7..3902b24b3 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -2652,6 +2652,11 @@ def power(x1: ArrayLike, x2: ArrayLike, /) -> Array:
            [nan, 27.,  1.]], dtype=float32)
   """
   check_arraylike("power", x1, x2)
+
+  # Must do __jax_array__ conversion prior to dtype check.
+  x1 = x1.__jax_array__() if hasattr(x1, "__jax_array__") else x1
+  x2 = x2.__jax_array__() if hasattr(x2, "__jax_array__") else x2
+
   check_no_float0s("power", x1, x2)
 
   # We apply special cases, both for algorithmic and autodiff reasons:


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index 63a8762cd..45847b6f0 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -427,8 +427,8 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.polysub, Float[5], Float[5]),
   NumPyAPI.sig(jnp.polyval, Float[5], Float[10]),
   NumPyAPI.sig(jnp.positive, Float[5]),
-  # NumPyAPI.sig(jnp.pow, Float[5], Float[5]),
-  # NumPyAPI.sig(jnp.power, Float[5], Float[5]),
+  NumPyAPI.sig(jnp.pow, Float[5], Float[5]),
+  NumPyAPI.sig(jnp.power, Float[5], Float[5]),
   NumPyAPI.sig(jnp.prod, Float[5]),
   NumPyAPI.sig(jnp.ptp, Float[5]),
   NumPyAPI.sig(jnp.put, Float[5], Int[()], Float[()], inplace=False),
Repository: jax-ml/jax
Pull request #27588 (2025-03-29 01:24:23)

Fix an edge-case in reshape sharding rule where the last splitting/merging dim was `1`.
Fix an edge-case in reshape sharding rule where the last splitting/merging dim was `1`.



Patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index c2f0876ce..fd956136c 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -6977,11 +6977,16 @@ def _split_on_one_axis(op_shape, new_sizes, name):
             ' the sharding of the output via the `sharding` argument of'
             f' jax.lax.reshape. Got operand.shape={op_shape} and {new_sizes=}')
       temp = [new_sizes[j]]
-      while math.prod(temp) != op_shape[i]:
+      next_j = j + 1
+      while (math.prod(temp) != op_shape[i] or
+             (next_j < len(new_sizes) and new_sizes[next_j] == 1)):
         if math.prod(temp) > op_shape[i]:
           return False, []
         j += 1
+        if j >= len(new_sizes):
+          return False, []
         temp.append(new_sizes[j])
+        next_j += 1
       out.append(temp)
     i += 1
     j += 1


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index b49ba19c7..0b2daee8c 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -5508,6 +5508,18 @@ class ShardingInTypesTest(jtu.JaxTestCase):
       ('4', (1, 4, 1, 6, 1), (1, 4, 6),
        P(None, 'x', None, None, None), P(None, 'x', None), False),
       ('5', (4, 6), (4, 6), P(None, 'x'), P(None, 'x'), False),
+      ('6', (1024, 4096), (1024, 2048, 2, 1, 1, 1, 1),
+       P('x', None), P('x', None, None, None, None, None, None), False),
+      ('7', (1024, 4096, 32), (1024, 2048, 2, 1, 1, 32),
+       P('x', None, None), P('x', None, None, None, None, None), False),
+      ('8', (1024, 4096), (1024, 1, 1, 4096),
+       P('x', None), P('x', None, None, None), False),
+      ('9', (1024, 4096), (1024, 1, 1, 4096),
+       P(None, 'x'), P(None, None, None, 'x'), False),
+      ('10', (1024, 2048, 2, 1, 1, 1), (1024, 4096),
+       P('x', None, None, None, None, None), P('x', None), False),
+      ('11', (1024, 2048, 2, 1, 1, 1), (1024, 4096),
+       P(None, 'x', None, None, None, None), P(None, 'x'), False),
   )
   @jtu.with_user_mesh((2,), ('x',))
   def test_reshape(self, src_shape, dst_shape, src_spec, dst_spec,
@@ -5519,6 +5531,8 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     @partial(jax.jit, static_argnums=1)
     def f(x, new_sharding):
       y = lax.reshape(x, dst_shape, out_sharding=new_sharding)
+      self.assertEqual(y.aval.sharding.spec, dst_spec)
+      self.assertEqual(y.shape, dst_shape)
       y = y * 2
       self.assertEqual(y.aval.sharding.spec, dst_spec)
       return y
Repository: jax-ml/jax
Pull request #27627 (2025-03-31 22:16:00)

__jax_array__: add support in jnp.reshape, jnp.transpose, jnp.matrix_transpose


Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 63edaed0a..7b900e090 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1203,8 +1203,8 @@ def transpose(a: ArrayLike, axes: Sequence[int] | None = None) -> Array:
     Array([[1, 3],
            [2, 4]], dtype=int32)
   """
-  util.check_arraylike("transpose", a)
-  axes_ = list(range(np.ndim(a))[::-1]) if axes is None else axes
+  a = util.ensure_arraylike("transpose", a)
+  axes_ = list(range(a.ndim)[::-1]) if axes is None else axes
   axes_ = [_canonicalize_axis(i, np.ndim(a)) for i in axes_]
   return lax.transpose(a, axes_)
 
@@ -1285,8 +1285,8 @@ def matrix_transpose(x: ArrayLike, /) -> Array:
            [[5, 7],
             [6, 8]]], dtype=int32)
   """
-  util.check_arraylike("matrix_transpose", x)
-  ndim = np.ndim(x)
+  x = util.ensure_arraylike("matrix_transpose", x)
+  ndim = x.ndim
   if ndim < 2:
     raise ValueError(f"x must be at least two-dimensional for matrix_transpose; got {ndim=}")
   axes = (*range(ndim - 2), ndim - 1, ndim - 2)


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index fae9129dd..63a8762cd 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -375,7 +375,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.logical_or, Int[5], Int[5]),
   NumPyAPI.sig(jnp.logical_xor, Int[5], Int[5]),
   NumPyAPI.sig(jnp.matmul, Float[5, 5], Float[5]),
-  # NumPyAPI.sig(jnp.matrix_transpose, Float[5, 6]),
+  NumPyAPI.sig(jnp.matrix_transpose, Float[5, 6]),
   NumPyAPI.sig(jnp.matvec, Float[5, 5], Float[5]),
   NumPyAPI.sig(jnp.max, Float[5]),
   NumPyAPI.sig(jnp.maximum, Float[5], Float[5]),
@@ -442,7 +442,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.reciprocal, Float[5]),
   NumPyAPI.sig(jnp.remainder, Float[5], Float[5]),
   # NumPyAPI.sig(jnp.repeat, Float[5], Int[5]),
-  # NumPyAPI.sig(jnp.reshape, Float[6], (2, 3)),
+  NumPyAPI.sig(jnp.reshape, Float[6], shape=(2, 3)),
   NumPyAPI.sig(jnp.resize, Float[6], new_shape=(2, 3)),
   NumPyAPI.sig(jnp.right_shift, Int[5], Int[5]),
   NumPyAPI.sig(jnp.rint, Float[5]),
@@ -481,7 +481,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.tensordot, Float[2, 3, 4], Float[3, 4, 5]),
   # NumPyAPI.sig(jnp.tile, Float[5], reps=(2,)),
   NumPyAPI.sig(jnp.trace, Float[5, 5]),
-  # NumPyAPI.sig(jnp.transpose, Float[5, 6]),
+  NumPyAPI.sig(jnp.transpose, Float[5, 6]),
   NumPyAPI.sig(jnp.trapezoid, Float[5]),
   NumPyAPI.sig(jnp.tril, Float[5, 6]),
   NumPyAPI.sig(jnp.tril_indices_from, Float[5, 6]),
Repository: jax-ml/jax
Pull request #27644 (2025-04-01 17:06:13)

cumulative reductions: support __jax_array__ on inputs


Patch:
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index 985b296bc..96b2782ed 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -30,7 +30,7 @@ from jax._src import core
 from jax._src import deprecations
 from jax._src import dtypes
 from jax._src.numpy.util import (
-    _broadcast_to, check_arraylike, _complex_elem_type,
+    _broadcast_to, check_arraylike, _complex_elem_type, ensure_arraylike,
     promote_dtypes_inexact, promote_dtypes_numeric, _where)
 from jax._src.lax import lax as lax_internal
 from jax._src.typing import Array, ArrayLike, DType, DTypeLike, DeprecatedArg
@@ -1992,7 +1992,7 @@ def _cumulative_reduction(
     fill_nan: bool = False, fill_value: ArrayLike = 0,
     promote_integers: bool = False) -> Array:
   """Helper function for implementing cumulative reductions."""
-  check_arraylike(name, a)
+  a = ensure_arraylike(name, a)
   if out is not None:
     raise NotImplementedError(f"The 'out' argument to jnp.{name} is not supported")
   dtypes.check_user_dtype_supported(dtype, name)
@@ -2242,8 +2242,7 @@ def cumulative_sum(
     Array([[ 0,  1,  3,  6],
            [ 0,  4,  9, 15]], dtype=int32)
   """
-  check_arraylike("cumulative_sum", x)
-  x = lax_internal.asarray(x)
+  x = ensure_arraylike("cumulative_sum", x)
   if x.ndim == 0:
     raise ValueError(
       "The input must be non-scalar to take a cumulative sum, however a "
@@ -2304,8 +2303,7 @@ def cumulative_prod(
     Array([[  1,   1,   2,   6],
            [  1,   4,  20, 120]], dtype=int32)
   """
-  check_arraylike("cumulative_prod", x)
-  x = lax_internal.asarray(x)
+  x = ensure_arraylike("cumulative_prod", x)
   if x.ndim == 0:
     raise ValueError(
       "The input must be non-scalar to take a cumulative product, however a "


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index 69e9e1609..8f5ea33b5 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -283,10 +283,10 @@ NUMPY_APIS = [
   # NumPyAPI.sig(np.count_nonzero, [float], [(10,)]),
   # NumPyAPI.sig(np.cov, [float], [(10,)]),
   # NumPyAPI.sig(np.cross, [float, float], [(3,), (3,)]),
-  # NumPyAPI.sig(np.cumprod, [float], [(10,)]),
-  # NumPyAPI.sig(np.cumsum, [float], [(10,)]),
-  # NumPyAPI.sig(np.cumulative_prod, [float], [(10,)]),
-  # NumPyAPI.sig(np.cumulative_sum, [float], [(10,)]),
+  NumPyAPI.sig(jnp.cumprod, Float[5]),
+  NumPyAPI.sig(jnp.cumsum, Float[5]),
+  NumPyAPI.sig(jnp.cumulative_prod, Float[5]),
+  NumPyAPI.sig(jnp.cumulative_sum, Float[5]),
   NumPyAPI.sig(jnp.deg2rad, Float[5]),
   NumPyAPI.sig(jnp.degrees, Float[5]),
   # NumPyAPI.sig(jnp.delete, Float[5], Int[()]),
Repository: jax-ml/jax
Pull request #27648 (2025-04-01 18:07:00)

jnp.einsum: add support for __jax_array__


Patch:
diff --git a/jax/_src/numpy/einsum.py b/jax/_src/numpy/einsum.py
index 9d745643b..21333a9e7 100644
--- a/jax/_src/numpy/einsum.py
+++ b/jax/_src/numpy/einsum.py
@@ -288,6 +288,10 @@ def einsum(
   spec = operands[0] if isinstance(operands[0], str) else None
   path_type = 'optimal' if optimize is True else Unoptimized() if optimize is False else optimize
 
+  # Extract __jax_array__ before passing to contract_path()
+  operands = tuple(op.__jax_array__() if hasattr(op, "__jax_array__") else op
+                   for op in operands)
+
   # Allow handling of shape polymorphism
   non_constant_dim_types = {
       type(d) for op in operands if not isinstance(op, str)


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index f62491a60..69e9e1609 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -548,6 +548,29 @@ class JaxArrayTests(jtu.JaxTestCase):
     self.assertEqual(result.shape, obj.shape)
     self.assertEqual(result.dtype, obj.dtype)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "subscript-form", "args": ("jk,k->j", Float[5, 3], Float[3])},
+      {"testcase_name": "index-form", "args": (Float[5, 3], (0, 1), Float[3], (1,), (0,))},
+  )
+  def test_einsum(self, args):
+    rng = jtu.rand_default(self.rng())
+    def make_arg(arg):
+      if isinstance(arg, jax.ShapeDtypeStruct):
+        return rng(arg.shape, arg.dtype)
+      return arg
+    args = jax.tree.map(make_arg, args)
+
+    def wrap_array(arg):
+      if isinstance(arg, (jax.Array, np.ndarray)):
+        return JaxArrayWrapper(arg)
+      return arg
+    wrapped_args = jax.tree.map(wrap_array, args)
+
+    expected = jnp.einsum(*args)
+    actual = jnp.einsum(*wrapped_args)
+
+    self.assertAllClose(actual, expected, atol=0, rtol=0)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: jax-ml/jax
Pull request #27650 (2025-04-01 20:13:52)

jnp.isinf & friends: support __jax_array__


Patch:
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index 60e10b3be..3fe63545e 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -36,7 +36,7 @@ from jax._src.numpy import error as jnp_error
 from jax._src.numpy import reductions
 from jax._src.numpy.ufunc_api import ufunc
 from jax._src.numpy.util import (
-   check_arraylike, promote_args, promote_args_inexact,
+   check_arraylike, ensure_arraylike, promote_args, promote_args_inexact,
    promote_args_numeric, promote_dtypes_inexact, promote_dtypes_numeric,
    promote_shapes, _where, check_no_float0s)
 from jax._src.util import set_module
@@ -3500,7 +3500,7 @@ def isinf(x: ArrayLike, /) -> Array:
     >>> jnp.isinf(x)
     Array([False,  True, False,  True, False], dtype=bool)
   """
-  check_arraylike("isinf", x)
+  x = ensure_arraylike("isinf", x)
   dtype = dtypes.dtype(x)
   if dtypes.issubdtype(dtype, np.floating):
     return lax.eq(lax.abs(x), _constant_like(x, np.inf))
@@ -3513,7 +3513,7 @@ def isinf(x: ArrayLike, /) -> Array:
     return lax.full_like(x, False, dtype=np.bool_)
 
 
-def _isposneginf(infinity: float, x: ArrayLike, out) -> Array:
+def _isposneginf(infinity: float, x: Array, out) -> Array:
   if out is not None:
     raise NotImplementedError("The 'out' argument to isneginf/isposinf is not supported.")
   dtype = dtypes.dtype(x)
@@ -3556,6 +3556,7 @@ def isposinf(x, /, out=None):
     >>> jnp.isposinf(x)
     Array([False, False,  True, False, False], dtype=bool)
   """
+  x = ensure_arraylike("isposinf", x)
   return _isposneginf(np.inf, x, out)
 
 
@@ -3590,6 +3591,7 @@ def isneginf(x, /, out=None):
     >>> jnp.isneginf(x)
     Array([ True, False, False, False, False], dtype=bool)
   """
+  x = ensure_arraylike("isneginf", x)
   return _isposneginf(-np.inf, x, out)
 
 
@@ -3624,7 +3626,7 @@ def isnan(x: ArrayLike, /) -> Array:
     >>> jnp.isnan(x)
     Array([False, False, False,  True], dtype=bool)
   """
-  check_arraylike("isnan", x)
+  x = ensure_arraylike("isnan", x)
   return lax.ne(x, x)
 
 


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index 8f5ea33b5..55089720f 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -353,8 +353,8 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.isin, Int[5], Int[10]),
   NumPyAPI.sig(jnp.isinf, Float[5]),
   NumPyAPI.sig(jnp.isnan, Float[5]),
-  # NumPyAPI.sig(jnp.isneginf, Float[5]),
-  # NumPyAPI.sig(jnp.isposinf, Float[5]),
+  NumPyAPI.sig(jnp.isneginf, Float[5]),
+  NumPyAPI.sig(jnp.isposinf, Float[5]),
   NumPyAPI.sig(jnp.isreal, Float[5]),
   NumPyAPI.sig(jnp.isrealobj, Float[5]),
   NumPyAPI.sig(jnp.isscalar, Float[()]),
Repository: jax-ml/jax
Pull request #27651 (2025-04-01 20:19:01)

[mutable-arrays] add vmap rule for mutable_array_p, very basic test


Patch:
diff --git a/jax/_src/interpreters/batching.py b/jax/_src/interpreters/batching.py
index 03c9a9510..a187d4251 100644
--- a/jax/_src/interpreters/batching.py
+++ b/jax/_src/interpreters/batching.py
@@ -1169,3 +1169,8 @@ def add_batched(batched_args, batch_dims):
     x = moveaxis(x, bdx, bdy)
     return add_jaxvals(x, y), bdy
 primitive_batchers[add_jaxvals_p] = add_batched
+
+
+### mutable arrays
+
+defvectorized(core.mutable_array_p)


Test patch:
diff --git a/tests/mutable_array_test.py b/tests/mutable_array_test.py
index 950bddf54..a51e1d784 100644
--- a/tests/mutable_array_test.py
+++ b/tests/mutable_array_test.py
@@ -239,6 +239,16 @@ class MutableArrayTest(jtu.JaxTestCase):
       x_ref = core.mutable_array(x)
       y = f(x_ref)
 
+  def test_vmap_basic(self):
+    @jax.vmap
+    def f(x):
+      x_ref = core.mutable_array(x)
+      x_ref[...] =  x_ref[...] * x_ref[...]
+      return x_ref[...]
+    xs = jnp.arange(4.)
+    ys = f(xs)
+    self.assertAllClose(ys, xs ** 2, check_dtypes=False)
+
 
 @jtu.with_config(jax_mutable_array_checks=True)
 class MutableArrayErrorsTest(jtu.JaxTestCase):
Repository: jax-ml/jax
Pull request #27661 (2025-04-02 04:43:02)

Relax the aval check in `select_hlo_lowering_opaque` to only check for shardings if they are not empty. The same thing happens in select_p's sharding rule
Relax the aval check in `select_hlo_lowering_opaque` to only check for shardings if they are not empty. The same thing happens in select_p's sharding rule



Patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index fd956136c..ac6054328 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -7357,7 +7357,11 @@ def _select_jvp(primals, tangents):
 def _select_hlo_lowering_opaque(ctx, which, *cases):
   avals_in = ctx.avals_in
   aval_out, = ctx.avals_out
-  assert all(aval_case == aval_out for aval_case in avals_in[1:])
+  assert all((aval_case.shape, aval_case.dtype) == (aval_out.shape, aval_out.dtype)
+             for aval_case in avals_in[1:])
+  assert all(
+      aval_case == aval_out for aval_case in avals_in[1:]
+      if not aval_case.sharding.mesh.empty and not aval_out.sharding.mesh.empty)
   select_lower = _select_hlo_lowering
 
   physical_aval_out = core.physical_aval(aval_out)


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index ee4a8cd3e..38f191302 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7365,6 +7365,26 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     out = h(arr)
     self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
 
+  def test_scan_with_random_key_inside_jit(self):
+    mesh = jtu.create_mesh((2,), ('x',))
+    sharding = NamedSharding(mesh, P(None, 'x'))
+
+    @jax.jit
+    def scan(xs):
+      def step(carry, x):
+        next_carry = jax.vmap(jax.random.fold_in)(carry, x)
+        next_carry = jnp.where(x % 2 == 0, carry, next_carry)
+        return next_carry, None
+      rng = jnp.broadcast_to(jax.random.key(0), xs.shape[1:])
+      rng, _ = jax.lax.scan(step, rng, xs)
+      return rng
+
+    xs = jnp.arange(8).reshape(2, 4)
+    scan(xs)
+
+    xs = jax.device_put(xs, sharding)
+    scan(xs)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):