nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: jax-ml/jax
Pull request #27677 (2025-04-02 20:06:07)

Fix custom_transpose when composed with custom_jvp and use_direct_linearize=True
The real problem is that the partial eval rule for `custom_transpose`

https://github.com/jax-ml/jax/blob/92f7aeab48f144ba059cac29406b267e4030fe31/jax/_src/interpreters/partial_eval.py#L411-L428

should trace `call` to a Jaxpr, but it doesn't. It shouldn't be too hard to fix the partial eval rule for `custom_transpose`, but maybe it's worth merging this until I get to that?


Patch:
diff --git a/jax/_src/custom_transpose.py b/jax/_src/custom_transpose.py
index 5e87fdb20..21e607b5b 100644
--- a/jax/_src/custom_transpose.py
+++ b/jax/_src/custom_transpose.py
@@ -177,15 +177,19 @@ class CustomTransposePrimitive(core.Primitive):
   # TODO(frostig,mattjj): consider keeping `call` as a named parameter
   # instead of following this "call primitive" convention.
   def get_bind_params(self, params):
-    assert 'call_jaxpr' in params
-    assert 'transpose_jaxpr_thunk' in params
-    new_params: dict[str, Any] = dict(params)
-    new_params['transpose'] = make_transpose_from_thunk(
-        new_params.pop('transpose_jaxpr_thunk'),
-        new_params['lin_tree'])
-    call_jaxpr: core.ClosedJaxpr = new_params.pop('call_jaxpr')
-    call = lu.wrap_init(core.jaxpr_as_fun(call_jaxpr),
-                        debug_info=call_jaxpr.jaxpr.debug_info)
+    if 'call_jaxpr' in params:
+      assert 'transpose_jaxpr_thunk' in params
+      new_params: dict[str, Any] = dict(params)
+      new_params['transpose'] = make_transpose_from_thunk(
+          new_params.pop('transpose_jaxpr_thunk'),
+          new_params['lin_tree'])
+      call_jaxpr: core.ClosedJaxpr = new_params.pop('call_jaxpr')
+      call = lu.wrap_init(core.jaxpr_as_fun(call_jaxpr),
+                          debug_info=call_jaxpr.jaxpr.debug_info)
+    else:
+      assert 'transpose' in params
+      new_params: dict[str, Any] = dict(params)
+      call = new_params.pop("call")
     return [call], new_params
 
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index 032c09910..83264f10e 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -10387,6 +10387,28 @@ class CustomTransposeTest(jtu.JaxTestCase):
     self.assertAllClose(f_(x), g_(x))
     self.assertAllClose(f_t(x), g_t(x))
 
+  def test_compose_custom_jvp(self):
+    @jax.custom_jvp
+    def f(x):
+      return jnp.sin(x)
+
+    @f.defjvp
+    def f_jvp(primals, tangents):
+      x, = primals
+      dx, = tangents
+      return f(x), g(x, dx)
+
+    @custom_transpose
+    def g(x, dx):
+      return jnp.cos(x) * dx
+
+    @g.def_transpose
+    def gt(x, t):
+      return jnp.cos(x) * t
+
+    with config.use_direct_linearize(True):
+      self.assertAllClose(jax.grad(f)(0.5), jnp.cos(0.5))
+
 
 class CustomDceTest(jtu.JaxTestCase):
Repository: jax-ml/jax
Pull request #27723 (2025-04-04 01:14:14)

Add `auto_axes`, `explicit_axes` and `manual_axes` properties to Mesh and AbstractMesh
Add `auto_axes`, `explicit_axes` and `manual_axes` properties to Mesh and AbstractMesh



Patch:
diff --git a/jax/_src/mesh.py b/jax/_src/mesh.py
index a8003e693..00859f9b3 100644
--- a/jax/_src/mesh.py
+++ b/jax/_src/mesh.py
@@ -174,6 +174,21 @@ class _BaseMesh:
   def _any_axis_explicit(self) -> bool:
     return any_axis_types_match(self._axis_types, AxisType.Explicit)
 
+  @functools.cached_property
+  def auto_axes(self):
+    return tuple(n for n, t in safe_zip(self.axis_names, self._axis_types)
+                 if t == AxisType.Auto)
+
+  @functools.cached_property
+  def explicit_axes(self):
+    return tuple(n for n, t in safe_zip(self.axis_names, self._axis_types)
+                 if t == AxisType.Explicit)
+
+  @functools.cached_property
+  def manual_axes(self):
+    return tuple(n for n, t in safe_zip(self.axis_names, self._axis_types)
+                 if t == AxisType.Manual)
+
   @functools.cached_property
   def _axis_types_dict(self):
     if not self.axis_names:


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 901ce9521..2bdc54607 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1391,6 +1391,9 @@ class ShardingTest(jtu.JaxTestCase):
     self.assertDictEqual(
         mesh._axis_types_dict, {AxisType.Auto: ('y',), AxisType.Explicit: ('x',),
                           AxisType.Manual: ('z',)})
+    self.assertEqual(mesh.explicit_axes, ('x',))
+    self.assertEqual(mesh.auto_axes, ('y',))
+    self.assertEqual(mesh.manual_axes, ('z',))
 
     mesh = jax.make_mesh((1, 1, 1), ('x', 'y', 'z'),
                          axis_types=(Explicit, Explicit, Manual))