Repository: google/jax
Pull request #21716 (2024-06-06 20:33:30)

[export] Rename in_shardings and out_shardings field.
We rename `in_shardings` to `in_shardings_hlo` to remove confusion with JAX's use of `in_shardings`.
We also rename `xla_compatible_in_sharding` to `in_shardings_jax` since we do not have a XLACompatibleSharding type anymore.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d47f59253..ceb7b4912 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,9 @@ Remember to align the itemized text with the first line of an item within a list
 * Deprecations
   * `jax.sharding.XLACompatibleSharding` is deprecated. Please use
     `jax.sharding.Sharding`.
+  * `jax.experimental.Exported.in_shardings` has been renamed as
+    `jax.experimental.Exported.in_shardings_hlo`. Same for `out_shardings`.
+    The old names will be removed after 3 months.
   * Removed a number of previously-deprecated APIs:
     * from {mod}`jax.core`: `non_negative_dim`, `DimSize`, `Shape`
     * from {mod}`jax.lax`: `tie_in`
@@ -34,6 +37,11 @@ Remember to align the itemized text with the first line of an item within a list
     some did unintentionally. Going forward, we recommend explicit use of
     {func}`jax.vmap` in such cases.
 
+* New Functionality
+  * Added {func}`jax.experimental.Exported.in_shardings_jax` to construct
+    shardings that can be used with the JAX APIs from the HloShardings
+    that are stored in the `Exported` objects.
+
 ## jaxlib 0.4.29
 
 * Bug fixes
diff --git a/jax/_src/api.py b/jax/_src/api.py
index aa4579475..a3d99eda3 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1811,10 +1811,6 @@ def _cpp_pmap(
 
   pmap_f = wraps(fun)(cpp_mapped_f)
 
-  @api_boundary
-  def lower(*args, **kwargs):
-    return specialize(*args, **kwargs).lower()
-
   @api_boundary
   def specialize(*args, **kwargs):
     lowering_parameters = kwargs.pop(
@@ -1823,34 +1819,32 @@ def _cpp_pmap(
         fun, in_axes, out_axes, static_broadcasted_tuple, donate_tuple,
         devices, backend, axis_size, args, kwargs)
     abstract_args = list(map(shaped_abstractify, p.flat_args))
-    closed_jaxpr, xc_backend, replicas, shards, pci = pxla.get_pmap_jaxpr(
-        p.flat_fun, backend, axis_name,
+    lower_callable = partial(
+        pxla.lower_parallel_callable, p.flat_fun, backend, axis_name,
         axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
         devices=p.devices,
         name=p.flat_fun.__name__,
         in_axes=p.in_axes_flat,
         out_axes_thunk=p.out_axes_thunk,
-        avals=abstract_args)
-    lower_callable = partial(
-        pxla.lower_parallel_callable, p.flat_fun, axis_name,
+        donated_invars=p.donated_invars,
+        is_explicit_global_axis_size=p.is_explicit_global_axis_size,
+        avals=abstract_args,
+        lowering_parameters=lowering_parameters)
+    jaxpr, _, _, _, _ = pxla.get_pmap_jaxpr(
+        p.flat_fun, backend, axis_name,
         axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
         devices=p.devices,
         name=p.flat_fun.__name__,
         in_axes=p.in_axes_flat,
-        donated_invars=p.donated_invars,
-        is_explicit_global_axis_size=p.is_explicit_global_axis_size,
-        avals=abstract_args,
-        lowering_parameters=lowering_parameters,
-        closed_jaxpr=closed_jaxpr,
-        backend=xc_backend,
-        replicas=replicas,
-        shards=shards,
-        pci=pci)
+        out_axes_thunk=p.out_axes_thunk,
+        avals=abstract_args)
     args_info = stages.make_args_info(p.in_tree, abstract_args, donate_tuple)
-    return stages.Specialized(closed_jaxpr, args_info, p.flat_fun.__name__,
+    return stages.Specialized(jaxpr, args_info, p.flat_fun.__name__,
                               p.out_tree(), lower_callable)
 
-  pmap_f.lower = lower
+  pmap_f.lower = _pmap_lower(
+      fun, axis_name, in_axes, out_axes, static_broadcasted_tuple, devices,
+      backend, axis_size, donate_tuple)
   pmap_f.specialize = specialize
 
   return pmap_f
@@ -1858,6 +1852,47 @@ def _cpp_pmap(
 _pmap_cache_clears = weakref.WeakSet()  # type: ignore
 
 
+def _pmap_lower(fun, axis_name, in_axes, out_axes, static_broadcasted_tuple,
+                devices, backend, axis_size, donate_tuple):  # noqa: F811
+  """Make a ``lower`` method for pmapped functions."""
+  # If the function we returned from ``pmap`` were a class instance,
+  # this might naturally be a method, with ``fun`` as a ``self`` and
+  # all the other arguments stored as attributes.
+  @api_boundary
+  def lower(*args, **kwargs) -> stages.Lowered:
+    """Lower a parallel-mapped form of this function for the given arguments.
+
+    A parallel-mapped and lowered function is staged out of Python and
+    translated to a compiler's input language, possibly in a
+    backend-dependent manner. It is ready for compilation but is not yet
+    compiled. It represents a function intended for SPMD execution on
+    multiple devices.
+
+    Returns:
+      A ``Lowered`` instance representing the post-map lowering.
+    """
+    lowering_parameters = kwargs.pop(
+        '_experimental_lowering_parameters', mlir.LoweringParameters())
+    p = _prepare_pmap(
+        fun, in_axes, out_axes, static_broadcasted_tuple, donate_tuple,
+        devices, backend, axis_size, args, kwargs)
+    abstract_args = list(map(shaped_abstractify, p.flat_args))
+    computation = pxla.lower_parallel_callable(
+        p.flat_fun, backend, axis_name,
+        axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
+        devices=p.devices,
+        name=p.flat_fun.__name__,
+        in_axes=p.in_axes_flat,
+        out_axes_thunk=p.out_axes_thunk,
+        donated_invars=p.donated_invars,
+        is_explicit_global_axis_size=p.is_explicit_global_axis_size,
+        avals=abstract_args,
+        lowering_parameters=lowering_parameters)
+    return stages.Lowered.from_flat_info(
+        computation, p.in_tree, abstract_args, donate_tuple, p.out_tree())
+
+  return lower
+
 def jvp(
     fun: Callable, primals, tangents, has_aux: bool = False
   ) -> tuple[Any, ...]:
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index 7ddac0e6f..686ee0ced 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -557,17 +557,11 @@ def parallel_callable(fun: lu.WrappedFun,
                       donated_invars: Sequence[bool],
                       is_explicit_global_axis_size: bool,
                       *avals):
-  closed_jaxpr, xc_backend, replicas, shards, pci = get_pmap_jaxpr(
-      fun, backend_name, axis_name,
-      axis_size=axis_size, global_axis_size=global_axis_size,
-      devices=devices, name=fun.__name__, in_axes=in_axes,
-      out_axes_thunk=out_axes_thunk, avals=avals)
   pmap_computation = lower_parallel_callable(
-      fun, axis_name, axis_size, global_axis_size, devices, name,
-      in_axes, donated_invars,
+      fun, backend_name, axis_name, axis_size, global_axis_size, devices, name,
+      in_axes, out_axes_thunk, donated_invars,
       is_explicit_global_axis_size, avals,
-      lowering_parameters=mlir.LoweringParameters(), closed_jaxpr=closed_jaxpr,
-      backend=xc_backend, replicas=replicas, shards=shards, pci=pci)
+      lowering_parameters=mlir.LoweringParameters())
   pmap_executable = pmap_computation.compile()
   return WeakRefList([pmap_executable.unsafe_call, pmap_executable.fingerprint])
 
@@ -699,22 +693,19 @@ def get_pmap_jaxpr(
 @profiler.annotate_function
 def lower_parallel_callable(
     fun: lu.WrappedFun,
+    backend_name: str | None,
     axis_name: core.AxisName,
     axis_size: int,
     global_axis_size: int,
     devices: Sequence[xc.Device] | None,
     name: str,
     in_axes: Iterable[int | None],
+    out_axes_thunk: Callable[[], Sequence[int | None]],
     donated_invars: Sequence[bool],
     is_explicit_global_axis_size: bool,
     avals: Sequence[core.AbstractValue],
     *,
-    lowering_parameters: mlir.LoweringParameters,
-    closed_jaxpr: core.ClosedJaxpr,
-    backend: xc.Client,
-    replicas: ReplicaInfo,
-    shards: ShardInfo,
-    pci: ParallelCallableInfo) -> PmapComputation:
+    lowering_parameters: mlir.LoweringParameters) -> PmapComputation:
   # Determine global_axis_size for use in AxisEnv.
   # TODO(mattjj,skyewm): revive this check (inner_pmap always False now)
   # if xb.process_count() > 1 and global_axis_size is None and inner_pmap:
@@ -725,6 +716,9 @@ def lower_parallel_callable(
         f"Specified axis_size {global_axis_size} doesn't match received "
         f"axis_size {axis_size}.")
 
+  closed_jaxpr, backend, replicas, shards, pci = get_pmap_jaxpr(
+    fun, backend_name, axis_name, axis_size, global_axis_size, devices, name,
+    in_axes, out_axes_thunk, avals)
   jaxpr = closed_jaxpr.jaxpr
 
   no_nested_sharding = False
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index a139227ad..4a3d8400a 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -477,18 +477,29 @@ def _make_jit_wrapper(jit_info: PjitInfo):
 
   @api_boundary
   def lower(*args, **kwargs):
-    specialized = specialize(*args, **kwargs)
+    lowering_parameters = kwargs.pop(
+        '_experimental_lowering_parameters', mlir.LoweringParameters())
+
+    (args_flat, params, in_tree, out_tree,
+     donated_invars, arg_names, _) = _infer_params(jit_info, args, kwargs)
     try:
-      return specialized.lower()
+      lowering = _resolve_and_lower(
+          args_flat, **params, lowering_parameters=lowering_parameters)
     except pxla.DeviceAssignmentMismatchError as e:
       fails, = e.args
+      api_name = 'jit' if params['resource_env'] is None else 'pjit'
       fun = jit_info.fun
       fun_name = getattr(fun, '__qualname__',
                          getattr(fun, '__name__', str(fun)))
       msg = _device_assignment_mismatch_error(
-          fun_name, fails, specialized._args_flat, 'jit', specialized._arg_names)
+          fun_name, fails, args_flat, api_name, arg_names)
       raise ValueError(msg) from None
 
+    donate_argnums = tuple(i for i, d in enumerate(donated_invars) if d)
+    jaxpr = params["jaxpr"]
+    return stages.Lowered.from_flat_info(
+        lowering, in_tree, jaxpr.in_avals, donate_argnums, out_tree)
+
   @api_boundary
   def eval_shape(*args, **kwargs):
     _, params, _, out_tree, _, _, _ = _infer_params(jit_info, args, kwargs)
@@ -503,8 +514,8 @@ def _make_jit_wrapper(jit_info: PjitInfo):
     lowering_parameters = kwargs.pop(
         '_experimental_lowering_parameters', mlir.LoweringParameters())
 
-    (args_flat, params, in_tree, out_tree, donated_invars,
-     arg_names, _) = _infer_params(jit_info, args, kwargs)
+    args_flat, params, in_tree, out_tree, donated_invars, _, _ = _infer_params(
+        jit_info, args, kwargs)
 
     donate_argnums = tuple(i for i, d in enumerate(donated_invars) if d)
     jaxpr = params['jaxpr']
@@ -512,7 +523,7 @@ def _make_jit_wrapper(jit_info: PjitInfo):
     lower_callable = partial(_resolve_and_lower, args_flat, **params,
                              lowering_parameters=lowering_parameters)
     return stages.Specialized(jaxpr, args_info, params["name"], out_tree,
-                              lower_callable, args_flat, arg_names)
+                              lower_callable)
 
   wrapped = _cpp_pjit(jit_info)
   wrapped.lower = lower
@@ -643,6 +654,7 @@ def _infer_params(jit_info, args, kwargs):
   args_flat = [*implicit_args, *explicit_args]
 
   num_states_in = sum(init_tree.num_leaves for init_tree, _, _ in attrs_tracked)
+  num_states_out = sum(end_tree.num_leaves for _,  end_tree, _ in attrs_tracked)
   num_extra_args = len(implicit_args) + num_states_in + len(consts)
   in_shardings_flat = (UNSPECIFIED,) * num_extra_args + in_shardings_flat
   in_layouts_flat = (None,) * num_extra_args + in_layouts_flat
diff --git a/jax/_src/stages.py b/jax/_src/stages.py
index 90ad765a9..ba71ca655 100644
--- a/jax/_src/stages.py
+++ b/jax/_src/stages.py
@@ -426,18 +426,15 @@ class CompiledCallParams(NamedTuple):
 
 
 class Specialized(Stage):
-  __slots__ = ["jaxpr", "args_info", "fun_name", "_out_tree", "_lower_callable",
-               "_args_flat", "_arg_names"]
+  __slots__ = ["jaxpr", "args_info", "fun_name", "_out_tree", "_lower_callable"]
 
   def __init__(self, jaxpr: core.ClosedJaxpr, args_info, fun_name, out_tree,
-               lower_callable, args_flat=None, arg_names=None):
+               lower_callable):
     self.jaxpr = jaxpr
     self.args_info = args_info
     self.fun_name = fun_name
     self._out_tree = out_tree
     self._lower_callable = lower_callable
-    self._args_flat = args_flat
-    self._arg_names = arg_names
 
   @property
   def out_info(self):
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
index 774d5c830..8b13e8d63 100644
--- a/jax/experimental/export/_export.py
+++ b/jax/experimental/export/_export.py
@@ -150,15 +150,15 @@ class Exported:
     out_avals: the flat tuple of output abstract values. May contain dimension
         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.
+    in_shardings_hlo: the flattened input shardings, a sequence 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
+        the mesh. See `in_shardings_jax` 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.
+    out_shardings_hlo: the flattened output shardings, a sequence 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
+        the mesh. See `out_shardings_jax` 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',
@@ -290,8 +290,8 @@ class Exported:
   out_tree: tree_util.PyTreeDef
   out_avals: tuple[core.AbstractValue, ...]
 
-  in_shardings: tuple[HloSharding | None, ...]
-  out_shardings: tuple[HloSharding | None, ...]
+  in_shardings_hlo: tuple[HloSharding | None, ...]
+  out_shardings_hlo: tuple[HloSharding | None, ...]
   nr_devices: int
   lowering_platforms: tuple[str, ...]
   ordered_effects: tuple[effects.Effect, ...]
@@ -313,12 +313,21 @@ 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(
+  # For backwards compatibility
+  # TODO(necula): remove after September 2024.
+  @property
+  def in_shardings(self):
+    return self.in_shardings_hlo
+  @property
+  def out_shardings(self):
+    return self.out_shardings_hlo
+
+  def in_shardings_jax(
       self,
       mesh: sharding.Mesh) -> Sequence[sharding.Sharding | None]:
-    """Creates Shardings corresponding to self.in_shardings.
+    """Creates Shardings corresponding to self.in_shardings_hlo.
 
-    The Exported object stores `in_shardings` as HloShardings, which are
+    The Exported object stores `in_shardings_hlo` as HloShardings, which are
     independent of a mesh or set of devices. This method constructs
     Sharding that can be used in JAX APIs such as `jax.jit` or
     `jax.device_put`.
@@ -329,7 +338,7 @@ class Exported:
     >>> 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
+    >>> exp.in_shardings_hlo
     ({devices=[8]<=[8]},)
 
     # Create a mesh for running the exported object
@@ -337,7 +346,7 @@ class Exported:
     >>>
     # 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])
+    ...     exp.in_shardings_jax(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]),
@@ -350,17 +359,17 @@ class Exported:
      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)
+                 for s in self.in_shardings_hlo)
 
-  def xla_compatible_out_shardings(
+  def out_shardings_jax(
       self,
       mesh: sharding.Mesh) -> Sequence[sharding.Sharding | None]:
-    """Creates Shardings corresponding to self.out_shardings.
+    """Creates Shardings corresponding to self.out_shardings_hlo.
 
-    See documentation for xla_compatible_in_shardings.
+    See documentation for in_shardings_jax.
     """
     return tuple(_hlo_sharding_to_xla_compatible_sharding(s, mesh)
-                 for s in self.out_shardings)
+                 for s in self.out_shardings_hlo)
 
   def has_vjp(self) -> bool:
     return self._get_vjp is not None
@@ -584,9 +593,9 @@ def _export_lowered(
     fun_vjp_jax, vjp_in_avals = _get_vjp_fun(fun_jax,
                                              in_tree=exp_primal.in_tree,
                                              in_avals=exp_primal.in_avals,
-                                             in_shardings=exp_primal.in_shardings,
+                                             in_shardings_hlo=exp_primal.in_shardings_hlo,
                                              out_avals=exp_primal.out_avals,
-                                             out_shardings=exp_primal.out_shardings,
+                                             out_shardings_hlo=exp_primal.out_shardings_hlo,
                                              device_assignment=device_assignment,
                                              apply_jit=True,
                                              flat_primal_fun=True)
@@ -600,8 +609,8 @@ def _export_lowered(
       out_tree=lowered.out_tree,
       in_avals=tuple(args_avals_flat),
       out_avals=tuple(out_avals_flat),
-      in_shardings=in_shardings,
-      out_shardings=out_shardings,
+      in_shardings_hlo=in_shardings,
+      out_shardings_hlo=out_shardings,
       nr_devices=nr_devices,
       lowering_platforms=lowering._platforms,
       ordered_effects=ordered_effects,
@@ -1013,8 +1022,8 @@ 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[HloSharding | None, ...],
-                 out_shardings: tuple[HloSharding | None, ...],
+                 in_shardings_hlo: tuple[HloSharding | None, ...],
+                 out_shardings_hlo: tuple[HloSharding | None, ...],
                  device_assignment: Sequence[sharding_impls.Device] | None,
                  apply_jit: bool,
                  flat_primal_fun: bool = False,
@@ -1047,10 +1056,10 @@ def _get_vjp_fun(primal_fun: Callable, *,
     assert device_assignment is not None
     vjp_in_shardings = tuple(
         _hlo_sharding_to_gspmd_sharding(s, device_assignment)
-        for s in itertools.chain(in_shardings, out_shardings))
+        for s in itertools.chain(in_shardings_hlo, out_shardings_hlo))
     vjp_out_shardings = tuple(
         _hlo_sharding_to_gspmd_sharding(s, device_assignment)
-        for s in in_shardings)
+        for s in in_shardings_hlo)
     return pjit.pjit(fun_vjp_jax,
                      in_shardings=vjp_in_shardings,
                      out_shardings=vjp_out_shardings), vjp_in_avals
@@ -1210,7 +1219,7 @@ def _call_exported_lowering(ctx: mlir.LoweringRuleContext, *args,
       err_msg = "the module was lowered for more than 1 device."
     elif (_check_module(submodule, disabled_checks=()) or
           any(s is not None and not s.is_replicated()
-              for s in exported.in_shardings + exported.out_shardings)):
+              for s in exported.in_shardings_hlo + exported.out_shardings_hlo)):
       err_msg = "the module contains non-replicated sharding annotations."
     if err_msg:
       raise NotImplementedError(
@@ -1222,7 +1231,7 @@ def _call_exported_lowering(ctx: mlir.LoweringRuleContext, *args,
   # Apply in_shardings
   args = tuple(
     wrap_with_sharding(ctx, x, x_aval, x_sharding)
-    for x, x_aval, x_sharding in zip(args, ctx.avals_in, exported.in_shardings))
+    for x, x_aval, x_sharding in zip(args, ctx.avals_in, exported.in_shardings_hlo))
   symtab = ir.SymbolTable(submodule.operation)
   # The called function may have been exported with polymorphic shapes and called
   # now with more refined shapes. We insert hlo.ConvertOp to ensure the module
@@ -1311,7 +1320,7 @@ def _call_exported_lowering(ctx: mlir.LoweringRuleContext, *args,
   # Apply out_shardings
   results = tuple(
     wrap_with_sharding(ctx, x, x_aval, x_sharding)
-    for x, x_aval, x_sharding in zip(results, ctx.avals_out, exported.out_shardings)
+    for x, x_aval, x_sharding in zip(results, ctx.avals_out, exported.out_shardings_hlo)
   )
   return results
 
diff --git a/jax/experimental/export/_serialization.py b/jax/experimental/export/_serialization.py
index 3d58eed98..04928a2b2 100644
--- a/jax/experimental/export/_serialization.py
+++ b/jax/experimental/export/_serialization.py
@@ -79,10 +79,10 @@ def _serialize_exported(
   out_tree = _serialize_pytreedef(builder, exp.out_tree)
   out_avals = _serialize_array(builder, _serialize_aval, exp.out_avals)
   in_shardings = _serialize_array(
-      builder, _serialize_sharding, exp.in_shardings
+      builder, _serialize_sharding, exp.in_shardings_hlo
   )
   out_shardings = _serialize_array(
-      builder, _serialize_sharding, exp.out_shardings
+      builder, _serialize_sharding, exp.out_shardings_hlo
   )
   ordered_effects = _serialize_array(
       builder, _serialize_effect, exp.ordered_effects
@@ -212,8 +212,8 @@ def _deserialize_exported(exp: ser_flatbuf.Exported) -> export.Exported:
       out_tree=out_tree,
       out_avals=out_avals,
       nr_devices=nr_devices,
-      in_shardings=in_shardings,
-      out_shardings=out_shardings,
+      in_shardings_hlo=in_shardings,
+      out_shardings_hlo=out_shardings,
       lowering_platforms=lowering_platforms,
       ordered_effects=ordered_effects,
       unordered_effects=unordered_effects,
diff --git a/jax/experimental/jax2tf/jax2tf.py b/jax/experimental/jax2tf/jax2tf.py
index b86bbf0dc..655814e58 100644
--- a/jax/experimental/jax2tf/jax2tf.py
+++ b/jax/experimental/jax2tf/jax2tf.py
@@ -538,9 +538,9 @@ class NativeSerializationImpl(SerializationImpl):
     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,
+                                in_shardings_hlo=self.exported.in_shardings_hlo,
                                 out_avals=self.exported.out_avals,
-                                out_shardings=self.exported.out_shardings,
+                                out_shardings_hlo=self.exported.out_shardings_hlo,
                                 device_assignment=self.device_assignment,
                                 apply_jit=True)
 
@@ -610,9 +610,9 @@ class GraphSerializationImpl(SerializationImpl):
     return _export._get_vjp_fun(self.fun_jax,
                                 in_tree=self.in_tree,
                                 in_avals=self.args_avals_flat,
-                                in_shardings=(None,) * len(self.args_avals_flat),
+                                in_shardings_hlo=(None,) * len(self.args_avals_flat),
                                 out_avals=self.outs_avals,
-                                out_shardings=(None,) * len(self.outs_avals),
+                                out_shardings_hlo=(None,) * len(self.outs_avals),
                                 device_assignment=None,  # Not used when apply_jit = False
                                 apply_jit=False)
 
@@ -916,7 +916,7 @@ def _run_exported_as_tf(args_flat_tf: Sequence[TfVal],
   # See b/255511660.
   kept_in_shardings = []
   for i in exported.module_kept_var_idx:
-    kept_in_shardings.append(exported.in_shardings[i])
+    kept_in_shardings.append(exported.in_shardings_hlo[i])
   args_flat_tf = tuple(
     map(partial(_shard_value,
                 skip_replicated_sharding=tf.executing_eagerly()),
@@ -933,7 +933,7 @@ def _run_exported_as_tf(args_flat_tf: Sequence[TfVal],
 
   res = list(map(partial(_shard_value,
                          skip_replicated_sharding=tf.executing_eagerly()),
-                 res, exported.out_shardings))
+                 res, exported.out_shardings_hlo))
   res = tuple(map(_convert_value, res, exported.out_avals))
   return res
 


Test patch:
diff --git a/jax/_src/internal_test_util/export_back_compat_test_util.py b/jax/_src/internal_test_util/export_back_compat_test_util.py
index 3df15b4bb..7235c0ce2 100644
--- a/jax/_src/internal_test_util/export_back_compat_test_util.py
+++ b/jax/_src/internal_test_util/export_back_compat_test_util.py
@@ -330,8 +330,8 @@ data_{datetime.date.today().strftime('%Y_%m_%d')} = dict(
         in_avals=tuple(in_avals),
         out_tree=out_tree,
         out_avals=tuple(out_avals),
-        in_shardings=(None,) * len(in_avals),
-        out_shardings=(None,) * len(out_avals),
+        in_shardings_hlo=(None,) * len(in_avals),
+        out_shardings_hlo=(None,) * len(out_avals),
         lowering_platforms=(data.platform,),
         ordered_effects=(),
         unordered_effects=(),
diff --git a/tests/export_test.py b/tests/export_test.py
index ed42e36d0..30d239269 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -972,7 +972,7 @@ class JaxExportTest(jtu.JaxTestCase):
     # 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)
+    run_input_shardings = exp.in_shardings_jax(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)