Repository: google/jax
Pull request #21606 (2024-06-03 17:40:25)

[export] Add a helper function to create XLACompatibleShardings for in_shardings
This method can be used to construct shardings appropriate for sharding the arguments
when calling an Exported given a mesh.


Patch:
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
index 6647aa431..00e72603a 100644
--- a/jax/experimental/export/_export.py
+++ b/jax/experimental/export/_export.py
@@ -61,9 +61,7 @@ DType = Any
 Shape = jax._src.core.Shape
 # The values of input and output sharding from the lowering.
 LoweringSharding = Union[sharding.XLACompatibleSharding, pxla.UnspecifiedValue]
-
-# None means unspecified sharding
-Sharding = Union[xla_client.HloSharding, None]
+HloSharding = xla_client.HloSharding
 
 # See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions
 # for a description of the different versions.
@@ -153,7 +151,15 @@ class Exported:
         expressions in the shapes, with dimension variables among those in
         `in_avals.
     in_shardings: the flattened input shardings, as long as `in_avals`.
+        `None` means unspecified sharding.
+        Note that these do not include the mesh or the actual devices used in
+        the mesh. See `xla_compatible_in_shardings` for a way to turn these
+        into sharding specification that can be used with JAX APIs.
     out_shardings: the flattened output shardings, as long as `out_avals`.
+        `None` means unspecified sharding.
+        Note that these do not include the mesh or the actual devices used in
+        the mesh. See `xla_compatible_out_shardings` for a way to turn these
+        into sharding specification that can be used with JAX APIs.
     nr_devices: the number of devices that the module has been lowered for.
     lowering_platforms: a tuple containing at least one of 'tpu', 'cpu',
         'cuda', 'rocm'. See below for the calling convention for when
@@ -168,7 +174,7 @@ class Exported:
         See more versioning details at https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions.
     module_kept_var_idx: the sorted indices of the arguments among `in_avals` that
         must be passed to the module. The other arguments have been dropped
-        because they are not used. Same length as `in_shardings`.
+        because they are not used.
     uses_shape_polymorphism: whether the `mlir_module_serialized` uses shape
         polymorphism. This may be because `in_avals` contains dimension
         variables, or due to inner calls of Exported modules that have
@@ -284,8 +290,8 @@ class Exported:
   out_tree: tree_util.PyTreeDef
   out_avals: tuple[core.AbstractValue, ...]
 
-  in_shardings: tuple[Sharding, ...]
-  out_shardings: tuple[Sharding, ...]
+  in_shardings: tuple[HloSharding | None, ...]
+  out_shardings: tuple[HloSharding | None, ...]
   nr_devices: int
   lowering_platforms: tuple[str, ...]
   ordered_effects: tuple[effects.Effect, ...]
@@ -307,6 +313,55 @@ class Exported:
     # do not want the entire serialized module to end up in locations.
     return f"Exported(fun_name={self.fun_name}, ...)"
 
+  def xla_compatible_in_shardings(
+      self,
+      mesh: sharding.Mesh) -> Sequence[sharding.XLACompatibleSharding | None]:
+    """Creates XLACompatibleShardings corresponding to self.in_shardings.
+
+    The Exported object stores `in_shardings` as HloShardings, which are
+    independent of a mesh or set of devices. This method constructs
+    XLACompatibleSharding that can be used in JAX APIs such as `jax.jit` or
+    `jax.device_put`.
+
+    Example usage:
+    >>> from jax.experimental import export
+    >>> exp_mesh = sharding.Mesh(jax.devices(), ("a",))
+    >>> exp = export.export(jax.jit(lambda x: jax.numpy.add(x, x),
+    ...     in_shardings=sharding.NamedSharding(exp_mesh, sharding.PartitionSpec("a")))
+    ...     )(np.arange(jax.device_count()))
+    >>> exp.in_shardings
+    ({devices=[8]<=[8]},)
+
+    # Create a mesh for running the exported object
+    >>> run_mesh = sharding.Mesh(jax.devices()[::-1], ("b",))
+    >>>
+    # Put the args and kwargs on the appropriate devices
+    >>> run_arg = jax.device_put(np.arange(jax.device_count()),
+    ...     exp.xla_compatible_in_shardings(run_mesh)[0])
+    >>> res = export.call(exp)(run_arg)
+    >>> res.addressable_shards
+    [Shard(device=CpuDevice(id=7), index=(slice(0, 1, None),), replica_id=0, data=[0]),
+     Shard(device=CpuDevice(id=6), index=(slice(1, 2, None),), replica_id=0, data=[2]),
+     Shard(device=CpuDevice(id=5), index=(slice(2, 3, None),), replica_id=0, data=[4]),
+     Shard(device=CpuDevice(id=4), index=(slice(3, 4, None),), replica_id=0, data=[6]),
+     Shard(device=CpuDevice(id=3), index=(slice(4, 5, None),), replica_id=0, data=[8]),
+     Shard(device=CpuDevice(id=2), index=(slice(5, 6, None),), replica_id=0, data=[10]),
+     Shard(device=CpuDevice(id=1), index=(slice(6, 7, None),), replica_id=0, data=[12]),
+     Shard(device=CpuDevice(id=0), index=(slice(7, 8, None),), replica_id=0, data=[14])]
+    """
+    return tuple(_hlo_sharding_to_xla_compatible_sharding(s, mesh)
+                 for s in self.in_shardings)
+
+  def xla_compatible_out_shardings(
+      self,
+      mesh: sharding.Mesh) -> Sequence[sharding.XLACompatibleSharding | None]:
+    """Creates XLACompatibleShardings corresponding to self.out_shardings.
+
+    See documentation for xla_compatible_in_shardings.
+    """
+    return tuple(_hlo_sharding_to_xla_compatible_sharding(s, mesh)
+                 for s in self.out_shardings)
+
   def has_vjp(self) -> bool:
     return self._get_vjp is not None
 
@@ -486,7 +541,7 @@ def _export_lowered(
 
   nr_devices = len(lowering.compile_args["device_assignment"])
   def export_sharding(s: LoweringSharding,
-                      aval: core.ShapedArray) -> Sharding:
+                      aval: core.ShapedArray) -> HloSharding | None:
     if sharding_impls.is_unspecified(s):
       return None
     return s._to_xla_hlo_sharding(aval.ndim)  # type: ignore[union-attr]
@@ -921,33 +976,28 @@ def expand_in_shardings(in_shardings: Sequence[LoweringSharding],
     all_in_shardings[idx] = in_s
   return tuple(all_in_shardings)
 
-# TODO(yashkatariya, necula): remove this function once we relax the checks
-# in the jit front-end.
-def canonical_shardings(
-    device_assignment: Sequence[jax.Device],
-    in_shardings: Sequence[Sharding],
-    out_shardings: Sequence[Sharding]
-    ) -> tuple[Sequence[sharding.XLACompatibleSharding | None],
-               Sequence[sharding.XLACompatibleSharding | None]]:
-  """Prepares canonical in_ and out_shardings for a pjit invocation.
-
-  Turns the HloSharding into XLACompatibleSharding.
-
-  Returns: a pair with the canonicalized input and output shardings.
-  """
-  def canonicalize(
-    ss: Sequence[Sharding]) -> Sequence[sharding.XLACompatibleSharding | None]:
-    return tuple(
-        sharding.GSPMDSharding(device_assignment, s) if s is not None else None
-        for s in ss)
-  return (canonicalize(in_shardings), canonicalize(out_shardings))
+def _hlo_sharding_to_xla_compatible_sharding(
+    hlo_sharding: HloSharding | None,
+    mesh: sharding.Mesh) -> sharding.XLACompatibleSharding | None:
+  if hlo_sharding is None:
+    return None
+  return sharding_impls._gspmd_to_named_sharding_via_mesh(
+      _hlo_sharding_to_gspmd_sharding(hlo_sharding, tuple(mesh.devices.flat)),  # type: ignore[arg-type]
+      mesh)
+
+def _hlo_sharding_to_gspmd_sharding(
+    hlo_sharding: HloSharding | None,
+    device_assignment: Sequence[jax.Device]) -> sharding.GSPMDSharding | None:
+  if hlo_sharding is None:
+    return None
+  return sharding.GSPMDSharding(device_assignment, hlo_sharding)
 
 def _get_vjp_fun(primal_fun: Callable, *,
                  in_tree: tree_util.PyTreeDef,
                  in_avals: Sequence[core.AbstractValue],
                  out_avals: Sequence[core.AbstractValue],
-                 in_shardings: tuple[Sharding, ...],
-                 out_shardings: tuple[Sharding, ...],
+                 in_shardings: tuple[HloSharding | None, ...],
+                 out_shardings: tuple[HloSharding | None, ...],
                  device_assignment: Sequence[sharding_impls.Device] | None,
                  apply_jit: bool,
                  flat_primal_fun: bool = False,
@@ -978,10 +1028,12 @@ def _get_vjp_fun(primal_fun: Callable, *,
 
   if apply_jit:
     assert device_assignment is not None
-    vjp_in_shardings, vjp_out_shardings = canonical_shardings(
-      device_assignment,
-      tuple(itertools.chain(in_shardings, out_shardings)),
-      in_shardings)
+    vjp_in_shardings = tuple(
+        _hlo_sharding_to_gspmd_sharding(s, device_assignment)
+        for s in itertools.chain(in_shardings, out_shardings))
+    vjp_out_shardings = tuple(
+        _hlo_sharding_to_gspmd_sharding(s, device_assignment)
+        for s in in_shardings)
     return pjit.pjit(fun_vjp_jax,
                      in_shardings=vjp_in_shardings,
                      out_shardings=vjp_out_shardings), vjp_in_avals
@@ -1251,7 +1303,7 @@ mlir.register_lowering(call_exported_p, _call_exported_lowering)
 def wrap_with_sharding(ctx: mlir.LoweringRuleContext,
                        x: ir.Value,
                        x_aval: core.AbstractValue,
-                       x_sharding: Sharding) -> ir.Value:
+                       x_sharding: HloSharding | None) -> ir.Value:
   if x_sharding is None:
     return x
   return mlir.wrap_with_sharding_op(
diff --git a/jax/experimental/export/_serialization.py b/jax/experimental/export/_serialization.py
index c9dddbce4..3d58eed98 100644
--- a/jax/experimental/export/_serialization.py
+++ b/jax/experimental/export/_serialization.py
@@ -14,6 +14,8 @@
 
 # Serialization and deserialization of export.Exported
 
+from __future__ import annotations
+
 from typing import Callable, TypeVar
 from collections.abc import Sequence
 from functools import partial
@@ -366,7 +368,7 @@ def _deserialize_aval(aval: ser_flatbuf.AbstractValue,
 
 
 def _serialize_sharding(
-    builder: flatbuffers.Builder, s: _export.Sharding
+    builder: flatbuffers.Builder, s: _export.HloSharding | None
 ) -> int:
   proto = None
   if s is None:
@@ -383,7 +385,7 @@ def _serialize_sharding(
   return ser_flatbuf.ShardingEnd(builder)
 
 
-def _deserialize_sharding(s: ser_flatbuf.Sharding) -> _export.Sharding:
+def _deserialize_sharding(s: ser_flatbuf.Sharding) -> _export.HloSharding | None:
   kind = s.Kind()
   if kind == ser_flatbuf.ShardingKind.unspecified:
     return None


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index 69d6aff38..461735dc1 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -346,7 +346,6 @@ jax_test(
 jax_test(
     name = "infeed_test",
     srcs = ["infeed_test.py"],
-    tags = ["test_cpu_thunks"],
     deps = [
         "//jax:experimental_host_callback",
     ],
@@ -356,7 +355,6 @@ jax_test(
     name = "jax_jit_test",
     srcs = ["jax_jit_test.py"],
     main = "jax_jit_test.py",
-    tags = ["test_cpu_thunks"],
 )
 
 py_test(
diff --git a/tests/export_test.py b/tests/export_test.py
index 8c0cfb415..917d86420 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -936,6 +936,31 @@ class JaxExportTest(jtu.JaxTestCase):
         in_shardings=(jax.sharding.NamedSharding(mesh1, P("x", None)),)
       )(a)
 
+  def test_input_shardings_unused_args(self):
+    nr_devices = 2
+    if len(jax.devices()) < nr_devices:
+      self.skipTest("Need at least 2 devices")
+    devices = jax.devices()[0:nr_devices]
+    export_mesh = Mesh(np.array(devices),
+                       axis_names=("x",))
+    a = np.arange(16 * 4, dtype=np.float32).reshape((16, 4))
+
+    f = jax.jit(lambda x, y: jnp.sin(x),
+                in_shardings=(jax.sharding.NamedSharding(export_mesh, P("x", None),),
+                              None),
+                out_shardings=(jax.sharding.NamedSharding(export_mesh, P("x", None),)))
+    exp = get_exported(f)(a, a)
+
+    # We can use other devices and other meshes for running
+    run_devices = devices[::-1]
+    run_mesh = Mesh(run_devices, "a")
+    run_input_shardings = exp.xla_compatible_in_shardings(run_mesh)
+    a_run = jax.device_put(a, run_input_shardings[0])
+    b_run = jax.device_put(a, run_input_shardings[1])
+    res = export.call(exp)(a_run, b_run)
+    self.assertEqual(res.addressable_shards[0].device, run_devices[0])
+    self.assertEqual(res.addressable_shards[1].device, run_devices[1])
+
   def test_call_with_different_no_of_devices(self):
     if jax.local_device_count() < 2:
       self.skipTest("Need at least 2 devices")