Repository: google/jax
Pull request #21448 (2024-05-28 00:31:34)

[export] Simplify construction of shardings for the VJP
This is possible now due to improvements in the handling of shardings in pjit.

At the same time, re-enable the checking of shardings for the arguments and results of the VJP function.


Patch:
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
index 82298a0ec..c90a80b1c 100644
--- a/jax/experimental/export/_export.py
+++ b/jax/experimental/export/_export.py
@@ -899,27 +899,18 @@ def canonical_shardings(
     device_assignment: Sequence[jax.Device],
     in_shardings: Sequence[Sharding],
     out_shardings: Sequence[Sharding]
-    ) -> tuple[(pxla.UnspecifiedValue |
-                     Sequence[sharding.XLACompatibleSharding]),
-               (pxla.UnspecifiedValue |
-                     Sequence[sharding.XLACompatibleSharding])]:
+    ) -> tuple[Sequence[sharding.XLACompatibleSharding | None],
+               Sequence[sharding.XLACompatibleSharding | None]]:
   """Prepares canonical in_ and out_shardings for a pjit invocation.
 
-  The pjit front-end is picky about what in- and out-shardings it accepts,
-  e.g., if all are unspecified then the whole sharding should be the
-  sharding_impls.UNSPECIFIED object, otherwise the unspecified shardings are
-  replaced with the replicated sharding.
+  Turns the HloSharding into XLACompatibleSharding.
 
   Returns: a pair with the canonicalized input and output shardings.
   """
-  replicated_s = sharding.GSPMDSharding.get_replicated(device_assignment)
   def canonicalize(
-    ss: Sequence[Sharding]) -> (pxla.UnspecifiedValue |
-                                     Sequence[sharding.XLACompatibleSharding]):
-    if all(s is None for s in ss):
-      return sharding_impls.UNSPECIFIED
+    ss: Sequence[Sharding]) -> Sequence[sharding.XLACompatibleSharding | None]:
     return tuple(
-        sharding.GSPMDSharding(device_assignment, s) if s is not None else replicated_s
+        sharding.GSPMDSharding(device_assignment, s) if s is not None else None
         for s in ss)
   return (canonicalize(in_shardings), canonicalize(out_shardings))
 


Test patch:
diff --git a/tests/export_test.py b/tests/export_test.py
index 472da235a..96dd1d6ac 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -1055,24 +1055,26 @@ class JaxExportTest(jtu.JaxTestCase):
   @jtu.parameterized_filterable(
     kwargs=[
       dict(in_shardings=in_shardings, out_shardings=out_shardings,
-           with_mesh=with_mesh)
+           with_mesh_context=with_mesh_context)
       for in_shardings in ("missing", None, "P")
       for out_shardings in ("missing", None, "P")
-      for with_mesh in (True, False)
+      for with_mesh_context in (True, False)
   ])
   def test_grad_with_sharding(self, in_shardings="P", out_shardings=None,
-                              with_mesh=False):
+                              with_mesh_context=False):
     if len(jax.devices()) < 2:
       self.skipTest("Test requires at least 2 devices")
     x_shape = (10, 20)
     x = np.arange(np.prod(x_shape), dtype=np.float32).reshape(x_shape)
+    # The input has shape f32[10,20] and output f32[20,10] in order to
+    # distinguish them in the HLO.
     def f_jax(x):  # x: f32[10,20] -> f32[20,10]
       return jnp.sin(x.T)
 
     mesh = Mesh(jax.devices()[:2], "d")
     pjit_kwargs = {}
     # Use NamedShardings if we don't have a mesh_context
-    if with_mesh:
+    if with_mesh_context:
       sharding_None_d = P(None, "d")
       sharding_d_None = P("d", None)
     else:
@@ -1088,7 +1090,7 @@ class JaxExportTest(jtu.JaxTestCase):
     f_jax_pjit = pjit.pjit(f_jax, **pjit_kwargs)
 
     with contextlib.ExitStack() as stack:
-      if with_mesh:
+      if with_mesh_context:
         stack.enter_context(mesh)
       # Serialize higher-order gradiends
       exp = get_exported(f_jax_pjit, vjp_order=2)(x)
@@ -1098,50 +1100,62 @@ class JaxExportTest(jtu.JaxTestCase):
 
     vjp_module_str = str(exp_vjp.mlir_module())
 
+    # The MHLO attributes of the args and the result of the main function
+    # Arg0 are the primal inputs, arg1 are the output cotangent, res is the input cotangent
+    arg0_attrs, arg1_attrs, res_attrs = re.search(
+        r"func.func public @main\(%arg0: tensor<10x20xf32> (.*)"
+        r", %arg1: tensor<20x10xf32> (.*)"
+        r"\) -> \(tensor<10x20xf32> (.*)",  # the result
+        vjp_module_str).groups()
+
     if in_shardings == "P":
+      self.assertRegex(arg0_attrs, re.escape("{devices=[1,2]<=[2]}"))
+      self.assertRegex(res_attrs, re.escape("{devices=[1,2]<=[2]}"))
       primal_in_sharding = "{devices=[1,2]<=[2]}"
     else:
       primal_in_sharding = "{replicated}"
+      if with_mesh_context:
+        self.assertRegex(arg0_attrs, re.escape("replicated"))
+        self.assertRegex(res_attrs, re.escape("replicated"))
+      else:
+        # If there is no mesh context, we have used NamedSharding(None)
+        # and then the sharding is unspecified!
+        self.assertNotIn("mhlo.sharding", arg0_attrs)
+        self.assertNotIn("mhlo.sharding", res_attrs)
+
     if out_shardings == "P":
+      self.assertRegex(arg1_attrs, re.escape("{devices=[2,1]<=[2]}"))
       primal_out_sharding = "{devices=[2,1]<=[2]}"
     else:
       primal_out_sharding = "{replicated}"
+      if with_mesh_context:
+        self.assertRegex(arg1_attrs, re.escape("replicated"))
+      else:
+        self.assertNotIn("mhlo.sharding", arg1_attrs)
 
-    # TODO(b/326476605): Change the condition below if required.
-    if in_shardings == "P":
-      main = re.compile(
-        r"func.func public @main\(%arg0: tensor<10x20xf32>.*"
-        "mhlo.sharding = \"" + re.escape(primal_in_sharding) + "\""
-        r".*%arg1: tensor<20x10xf32>.*"
-        "mhlo.sharding = \"" + re.escape(primal_out_sharding) + "\""
-        # result
-        r".*->.*\(tensor<10x20xf32>.*"
-        "mhlo.sharding = \"" + re.escape(primal_in_sharding) + "\"")
-      self.assertRegex(vjp_module_str, main)
-
-    # Custom calls for the primal input shape all match primal_in_sharding
-    primal_in_calls = re.findall(
+    # Sharding custom calls for the primal input shape all match primal_in_sharding
+    primal_in_sharding_calls = re.findall(
       r"custom_call @Sharding.*mhlo.sharding = \"(.+)\".*:.*tensor<10x20xf32>",
       vjp_module_str)
     self.assertTrue(
-      all(s == primal_in_sharding for s in primal_in_calls),
-      primal_in_calls
+      all(s == primal_in_sharding for s in primal_in_sharding_calls),
+      primal_in_sharding_calls
     )
 
     # Custom calls for the primal output shape all match primal_out_sharding
-    primal_out_calls = re.findall(
+    primal_out_sharding_calls = re.findall(
       r"custom_call @Sharding.*mhlo.sharding = \"(.+)\".*:.*tensor<20x10xf32>",
       vjp_module_str)
     self.assertTrue(
-      all(s == primal_out_sharding for s in primal_out_calls),
-      primal_in_calls
+      all(s == primal_out_sharding for s in primal_out_sharding_calls),
+      primal_out_sharding_calls
     )
 
     # Call the exported gradient functions. In order to set the device context
     # we replicate the inputs. If we don't use a mesh context and there are
     # no shardings on inputs or outputs, then we have serialized for one
     # device.
-    if in_shardings != "P" and out_shardings != "P" and not with_mesh:
+    if in_shardings != "P" and out_shardings != "P" and not with_mesh_context:
       self.assertEqual(exp_vjp.nr_devices, 1)
       self.assertEqual(exp_vjp2.nr_devices, 1)
       call_mesh = Mesh(jax.devices()[:1], "e")