nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: jax-ml/jax
Pull request #27766 (2025-04-07 05:56:05)

Raise an error if the type passed to `axis_types` argument of `Mesh` and `AbstractMesh` is not `jax.sharding.AxisType`.
Raise an error if the type passed to `axis_types` argument of `Mesh` and `AbstractMesh` is not `jax.sharding.AxisType`.



Patch:
diff --git a/jax/_src/mesh.py b/jax/_src/mesh.py
index 00859f9b3..8db444554 100644
--- a/jax/_src/mesh.py
+++ b/jax/_src/mesh.py
@@ -111,12 +111,16 @@ class AxisType(enum.Enum):
   def __repr__(self):
     return self.name
 
-def _normalize_axis_types(axis_names, axis_types):
+def _normalize_axis_types(axis_names, axis_types, name):
   axis_types = ((AxisType.Auto,) * len(axis_names)
                 if axis_types is None else axis_types)
   if not isinstance(axis_types, tuple):
-    assert isinstance(axis_types, AxisType), axis_types
     axis_types = (axis_types,)
+
+  if not all(isinstance(a, AxisType) for a in axis_types):
+    raise TypeError(
+        f"axis_types passed to {name} must be of type `jax.sharding.AxisType`."
+        f" Got {axis_types} of type {tuple(type(a) for a in axis_types)}")
   if len(axis_names) != len(axis_types):
     raise ValueError(
         "Number of axis names should match the number of axis_types. Got"
@@ -256,7 +260,7 @@ class Mesh(_BaseMesh, contextlib.ContextDecorator):
           f"devices.ndim == {devices.ndim} and "
           f"len(axis_names) == {len(axis_names)}.")
 
-    axis_types = _normalize_axis_types(axis_names, axis_types)
+    axis_types = _normalize_axis_types(axis_names, axis_types, 'Mesh')
 
     key = (axis_names, devices.shape, tuple(devices.flat), axis_types)
     val = _mesh_object_dict.get(key, None)
@@ -440,7 +444,8 @@ class AbstractMesh(_BaseMesh):
     self.axis_sizes = axis_sizes
     self.axis_names = axis_names
     self._size = math.prod(self.axis_sizes) if self.axis_sizes else 0
-    self._axis_types = _normalize_axis_types(self.axis_names, axis_types)
+    self._axis_types = _normalize_axis_types(
+        self.axis_names, axis_types, 'AbstractMesh')
     self._hash = hash((self.axis_sizes, self.axis_names, self._axis_types))
 
   def __hash__(self):


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 76aa1093e..f097497ce 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1377,6 +1377,16 @@ class ShardingTest(jtu.JaxTestCase):
       jax.sharding.AbstractMesh((2, 1), ('x', 'y'),
                                 axis_types=jax.sharding.AxisType.Auto)
 
+    with self.assertRaisesRegex(TypeError, "axis_types.*must be of type"):
+      AbstractMesh((2,), ('x',), axis_types=("explicit",))
+
+    with self.assertRaisesRegex(TypeError, "axis_types.*must be of type"):
+      AbstractMesh((2,), ('x',), axis_types="explicit")
+
+    with self.assertRaisesRegex(TypeError, "axis_types.*must be of type"):
+      AbstractMesh((2, 2), ('x', 'y'),
+                   axis_types=("explicit", AxisType.Explicit))
+
   def test_make_mesh_axis_types(self):
     Auto, Explicit, Manual = AxisType.Auto, AxisType.Explicit, AxisType.Manual
Repository: jax-ml/jax
Pull request #27843 (2025-04-08 19:46:38)

Add JVP rule for linear_call
The JVP of a linear function is trivial, and I think it's reasonable to expect `jax.custom_derivatives.linear_call` to be differentiable wrt the primal inputs. This is useful because `linear_call` is our only pseudo-public API for labeling traceables as linear, a request that comes up occasionally.

#jax-fixit 


Patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index d0eb4c525..a8f136477 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -1447,6 +1447,19 @@ def _linear_call_impl(*args, callee, transpose_thunk, num_callee_consts,
   del transpose_thunk, num_callee_consts, num_res
   return core.eval_jaxpr(callee.jaxpr, (), *args)
 
+def _linear_call_jvp_rule(primals, tangents, callee, transpose_thunk,
+                          num_callee_consts, num_res):
+  consts_and_res, primals = split_list(primals, [num_callee_consts + num_res])
+  const_tangents, tangents = split_list(tangents, [num_callee_consts + num_res])
+  assert all(type(t) is Zero for t in const_tangents)
+  primals_out = linear_call_p.bind(
+      *consts_and_res, *primals, callee=callee, transpose_thunk=transpose_thunk,
+      num_callee_consts=num_callee_consts, num_res=num_res)
+  tangents_out = linear_call_p.bind(
+      *consts_and_res, *tangents, callee=callee, transpose_thunk=transpose_thunk,
+      num_callee_consts=num_callee_consts, num_res=num_res)
+  return primals_out, tangents_out
+
 def _linear_call_transpose_rule(cts, *args, callee, transpose_thunk,
                                 num_callee_consts, num_res):
   transpose, t_consts = transpose_thunk()
@@ -1478,6 +1491,7 @@ linear_call_p = core.Primitive('linear_call')
 linear_call_p.multiple_results = True
 linear_call_p.def_impl(_linear_call_impl)
 linear_call_p.def_abstract_eval(_linear_call_abstract_eval)
+ad.primitive_jvps[linear_call_p] = _linear_call_jvp_rule
 ad.primitive_transposes[linear_call_p] = _linear_call_transpose_rule
 xla.register_initial_style_primitive(linear_call_p)
 mlir.register_lowering(linear_call_p, mlir.lower_fun(


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index c07748806..ec946003d 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -10073,6 +10073,19 @@ class CustomTransposeTest(jtu.JaxTestCase):
       return jax.custom_derivatives.linear_call(fn, tp, None, x)
     jax.jit(f)(0.1)
 
+  def test_linear_call_grad(self):
+    def f(x, y):
+      def fn(r, x): return x / r
+      def tp(r, t): return t / r
+      return x + jax.custom_derivatives.linear_call(fn, tp, y, x)
+
+    def f_ref(x, y):
+      return x + x / y
+
+    x = jnp.array(6.)
+    y = jnp.array(3.)
+    self.assertAllClose(jax.grad(f)(x, y), jax.grad(f_ref)(x, y))
+
   def test_basic(self):
     def f(x, y):
       @custom_transpose(jnp.ones(2))
Repository: jax-ml/jax
Pull request #28071 (2025-04-16 20:47:07)

Fix a bug in reduce_window sharding rule where padding is a tuple but we were checking for a scalar instead. Fixes https://github.com/jax-ml/jax/issues/28070
Fix a bug in reduce_window sharding rule where padding is a tuple but we were checking for a scalar instead. Fixes https://github.com/jax-ml/jax/issues/28070



Patch:
diff --git a/jax/_src/lax/windowed_reductions.py b/jax/_src/lax/windowed_reductions.py
index 8ebdcd6c3..acbfae37e 100644
--- a/jax/_src/lax/windowed_reductions.py
+++ b/jax/_src/lax/windowed_reductions.py
@@ -530,7 +530,7 @@ def reduce_window_sharding_rule(operand, window_dimensions, window_strides,
       base_dilation, window_dilation):
     if spec is None:
       continue
-    if not (wdim == 1 and ws == 1 and pd == 1 and bd == 1 and wdil == 1):
+    if not (wdim == 1 and ws == 1 and pd == (0, 0) and bd == 1 and wdil == 1):
       raise core.ShardingTypeError(
           "Only trivial windowing is supported along non-replicated"
           f" dimensions. Got {operand.sharding.spec=}")
@@ -639,7 +639,8 @@ def _reduce_window_lower(
 ):
 
   operand_aval, = ctx.avals_in
-  scalar_aval = operand_aval.update(shape=())
+  scalar_aval = operand_aval.update(
+      shape=(), sharding=operand_aval.sharding.with_spec(()))
 
   return mlir.reduce_window(
       ctx,


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 025512121..8f30475ee 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7034,6 +7034,18 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     self.assertArraysEqual(out, np.cumsum(np_inp))
     self.assertEqual(out.sharding, NamedSharding(mesh, P(None)))
 
+    @jax.jit
+    def f(x):
+      x = jnp.expand_dims(x, 1)
+      self.assertEqual(x.aval.sharding.spec, P('x', None))
+      out = jnp.cumsum(x, axis=1)
+      self.assertEqual(out.aval.sharding.spec, P('x', None))
+      return out
+
+    arr2 = jax.device_put(np.arange(8), P('x'))
+    out = f(arr2)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
+
   def test_device_put_under_use_mesh(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'))
     x = jnp.zeros((4, 4), dtype=jnp.int32)