Merge pull request #21479 from gnecula:jax2tf_fix_grad_mesh

PiperOrigin-RevId: 638328377

Author: jax authors

jax/experimental/export/_export.py

@@ -351,6 +351,7 @@ def export(fun_jax: Callable,
            *,
            lowering_platforms: Sequence[str] | None = None,
            disabled_checks: Sequence[DisabledSafetyCheck] = (),
+           _device_assignment_for_internal_jax2tf_use_only = None,
            ) -> Callable[..., Exported]:
   """Exports native serialization for a JAX function.
 
@@ -413,12 +414,15 @@ def do_export(*args_specs, **kwargs_specs) -> Exported:
         _experimental_lowering_parameters=mlir.LoweringParameters(
           platforms=actual_lowering_platforms,
         ))
-    return _export_lowered(lowered, disabled_checks=disabled_checks)
+    return _export_lowered(
+        lowered, disabled_checks=disabled_checks,
+        _device_assignment_for_internal_jax2tf_use_only=_device_assignment_for_internal_jax2tf_use_only)
   return do_export
 
 def _export_lowered(
     lowered: stages.Lowered,
     disabled_checks: Sequence[DisabledSafetyCheck] = (),
+    _device_assignment_for_internal_jax2tf_use_only = None,
   ) -> Exported:
   version = config.jax_serialization_version.value
   if (version < minimum_supported_serialization_version or
@@ -498,6 +502,8 @@ def export_sharding(s: LoweringSharding,
     for s, aval in zip(lowering.compile_args["out_shardings"], out_avals_flat))
 
   device_assignment = lowering.compile_args["device_assignment"]
+  if _device_assignment_for_internal_jax2tf_use_only is not None:
+    _device_assignment_for_internal_jax2tf_use_only[0] = device_assignment
   def _get_exported_vjp(exp_primal: Exported) -> Exported:
     # Turn the primal jaxpr into a function, in preparation for exporting
     # the VJP. Note that jaxpr_as_fun produces a function with flat arguments

jax/experimental/jax2tf/jax2tf.py

@@ -514,11 +514,15 @@ def _restore_context():
       _thread_local_state.call_tf_concrete_function_list = _prev_func_list
 
     self._restore_context = _restore_context
+    _exported_device_assignment = [None]
     self.exported = export.export(
         self.fun_jax,
         lowering_platforms=self.native_serialization_platforms,
-        disabled_checks=self.native_serialization_disabled_checks
+        disabled_checks=self.native_serialization_disabled_checks,
+        _device_assignment_for_internal_jax2tf_use_only=_exported_device_assignment,
     )(*self.args_specs, **self.kwargs_specs)
+    assert(_exported_device_assignment[0] is not None)
+    self.device_assignment = _exported_device_assignment[0]
 
   def after_conversion(self):
     self._restore_context()
@@ -531,15 +535,13 @@ def run_fun_tf(self,
 
   def get_vjp_fun(self) -> tuple[Callable,
                                  Sequence[core.AbstractValue]]:
-    # TODO(necula): use the actual device assignment from the primal function
-    device_assignment = jax.devices(jax.default_backend())[:self.exported.nr_devices]
     return _export._get_vjp_fun(self.fun_jax,
                                 in_tree=self.exported.in_tree,
                                 in_avals=self.exported.in_avals,
                                 in_shardings=self.exported.in_shardings,
                                 out_avals=self.exported.out_avals,
                                 out_shardings=self.exported.out_shardings,
-                                device_assignment=device_assignment,
+                                device_assignment=self.device_assignment,
                                 apply_jit=True)
 
 class GraphSerializationImpl(SerializationImpl):

jax/experimental/jax2tf/tests/sharding_test.py

@@ -455,22 +455,26 @@ def f_grad_tf(x_v, res_ct):
         ])
 
   def test_grad_sharding_different_mesh(self):
-    self.skipTest("TODO: fix the plumbing of device_assignment for jax2tf: https://github.com/google/jax/pull/21319")
     # Convert with two similar meshes, the only difference being
     # the order of the devices. grad should not fail.
     # https://github.com/google/jax/issues/21314
+    devices = jax.local_devices()[:2]
+    if len(devices) < 2:
+      raise unittest.SkipTest("Test requires 2 local devices")
     def f_jax(x):
       return jnp.sum(x * 2.)
 
-    mesh = Mesh(jax.local_devices(), "i")
+    mesh = Mesh(devices, "i")
     # The same mesh with reversed order of devices
-    mesh_rev = Mesh(list(reversed(jax.local_devices())), "i")
+    mesh_rev = Mesh(list(reversed(devices)), "i")
     shardings = NamedSharding(mesh, jax.sharding.PartitionSpec(("i",)))
     shardings_rev = NamedSharding(mesh_rev, jax.sharding.PartitionSpec(("i",)))
 
-    f_tf = tf.function(jax2tf.convert(pjit.pjit(f_jax, in_shardings=shardings)))
-    f_tf_rev = tf.function(jax2tf.convert(pjit.pjit(f_jax, in_shardings=shardings_rev)))
-    inp = np.ones((jax.local_device_count(), 4), dtype=np.float32)
+    f_tf = tf.function(jax2tf.convert(pjit.pjit(f_jax, in_shardings=shardings)),
+                       autograph=False)
+    f_tf_rev = tf.function(jax2tf.convert(pjit.pjit(f_jax, in_shardings=shardings_rev)),
+                           autograph=False)
+    inp = np.ones((2, 4), dtype=np.float32)
 
     input_v = tf.Variable(inp)
     with tf.GradientTape(persistent=True) as tape: