Remove the sharding and layout checks for non-DCE'd arguments during AOT safe call.

This is because the tracing, lowering and compilation caches do not register a miss if sharding/layout of a DCE'd arg changes when it's passed again to a jitted function.

This is not true for avals so that check still exists.

PiperOrigin-RevId: 623375760

Author: yashk2810

jax/_src/interpreters/pxla.py

@@ -1989,10 +1989,8 @@ def are_all_shardings_default_mem_kind(da_object, shardings):
 
 
 class AllArgsInfo(NamedTuple):
-  """Avals, shardings, layouts and debug_info for all arguments prior to DCE."""
+  """Avals and debug_info for all arguments prior to DCE."""
   in_avals: Sequence[core.ShapedArray]
-  in_shardings: Any
-  in_layouts: Any
   debug_info: core.JaxprDebugInfo | None
 
 
@@ -2038,8 +2036,7 @@ def lower_sharding_computation(
   auto_spmd_lowering = check_if_any_auto(
       it.chain.from_iterable([in_shardings, out_shardings]))  # type: ignore
 
-  all_args_info = AllArgsInfo(global_in_avals, in_shardings, in_layouts,
-                              closed_jaxpr.jaxpr.debug_info)
+  all_args_info = AllArgsInfo(global_in_avals, closed_jaxpr.jaxpr.debug_info)
 
   (closed_jaxpr, global_in_avals, global_out_avals, donated_invars,
    kept_var_idx, name_stack) = _dce_jaxpr(
@@ -3013,28 +3010,22 @@ def xla_extension_executable(self):
     return self.xla_executable
 
   def call(self, *args):
+    args_after_dce = [a for i, a in enumerate(args) if i in self._kept_var_idx]
     if self._all_args_info is None:
-      kept_args = [a for i, a in enumerate(args) if i in self._kept_var_idx]
+      kept_args = args_after_dce
       ref_avals = self.in_avals
-      in_shardings = self._in_shardings
-      in_layouts = self._in_layouts
       debug_info = None
     else:
       kept_args = args
       ref_avals = self._all_args_info.in_avals
-      iter_in_shardings = iter(self._in_shardings)
-      in_shardings = [next(iter_in_shardings) if i in self._kept_var_idx else s
-                      for i, s in enumerate(self._all_args_info.in_shardings)]
-      iter_in_layouts = iter(self._in_layouts)
-      in_layouts = [next(iter_in_layouts) if i in self._kept_var_idx else s
-                    for i, s in enumerate(self._all_args_info.in_layouts)]
       debug_info = self._all_args_info.debug_info
 
-    arg_avals = map(xla.abstractify, kept_args)
-    check_arg_avals_for_call(ref_avals, arg_avals, debug_info)
+    all_arg_avals = map(xla.abstractify, kept_args)
+    check_arg_avals_for_call(ref_avals, all_arg_avals, debug_info)
     # Check the GDA sharding and the input sharding.
-    check_array_xla_sharding_layout_match(kept_args, in_shardings,
-                                          in_layouts, debug_info)
+    check_array_xla_sharding_layout_match(
+        args_after_dce, self._in_shardings, self._in_layouts, debug_info,
+        self._kept_var_idx)
     return self.unsafe_call(*args)  # pylint: disable=not-callable
 
   def input_shardings(self) -> Sequence[sharding_impls.XLACompatibleSharding]:
@@ -3163,16 +3154,22 @@ def check_device_backend_on_shardings(shardings) -> bool:
 
 
 def check_array_xla_sharding_layout_match(
-    args, in_xla_shardings: Sequence[sharding_impls.XLACompatibleSharding],
+    args_after_dce,
+    in_xla_shardings: Sequence[sharding_impls.XLACompatibleSharding],
     in_xla_layouts: Sequence[DeviceLocalLayout],
-    jaxpr_debug_info: core.JaxprDebugInfo | None) -> None:
+    jaxpr_debug_info: core.JaxprDebugInfo | None,
+    kept_var_idx: set[int]) -> None:
   from jax._src.array import ArrayImpl
-  arg_names = ([''] * len(args) if jaxpr_debug_info is None else
-               jaxpr_debug_info.arg_names)
+  # jaxpr_debug_info.arg_names are before DCE, so need to DCE them.
+  arg_names = (
+      [""] * len(args_after_dce) if jaxpr_debug_info is None
+      else [a for i, a in enumerate(jaxpr_debug_info.arg_names)  # type: ignore
+            if i in kept_var_idx]
+  )
   errors = []
   num_errors = 5
-  for arg, xs, xl, name in safe_zip(args, in_xla_shardings, in_xla_layouts,
-                                    arg_names):
+  for arg, xs, xl, name in safe_zip(
+      args_after_dce, in_xla_shardings, in_xla_layouts, arg_names):
     if not isinstance(arg, ArrayImpl):
       continue
     if is_unspecified_or_auto(xs):
@@ -3200,7 +3197,6 @@ def check_array_xla_sharding_layout_match(
 
     if (xla_extension_version >= 249 and not db_xs and arg._committed and
         arg.layout.device_local_layout is not None and xl is not None and
-        not isinstance(xl, AutoLayout) and
         arg.layout.device_local_layout != xl):
       errors.append(
           ("Got input layout(s) that compiled object was called with: "

jax/_src/pjit.py

@@ -1486,7 +1486,7 @@ def _pjit_call_impl_python(
   if compiled._auto_spmd_lowering and config.enable_checks.value:
     pxla.check_array_xla_sharding_layout_match(
         args, compiled._in_shardings, compiled._in_layouts,
-        jaxpr.jaxpr.debug_info)
+        jaxpr.jaxpr.debug_info, compiled._kept_var_idx)
   if config.distributed_debug.value:
     # Defensively only perform fingerprint logic if debug logging is enabled
     # NOTE(skyewm): I didn't benchmark this

tests/pjit_test.py

@@ -4368,32 +4368,6 @@ def f(x, y):
         ' compiled'):
       g(x, y2)
 
-  def test_aot_error_on_dced_shardings_mismatch(self):
-    mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
-    shape = (8, 2)
-    np_inp = np.arange(math.prod(shape)).reshape(shape)
-
-    x = jax.device_put(np_inp, NamedSharding(mesh, P('x', 'y')))
-    y1 = jax.device_put(np_inp, NamedSharding(mesh, P('x')))
-    y2 = jax.device_put(np_inp, NamedSharding(mesh, P('y')))
-
-    @jax.jit
-    def f(x, y):
-      return x + 1
-
-    f_out1 = f(x, y1)
-    f(x, y2)
-
-    g = f.lower(x, y1).compile()
-    g_out1 = g(x, y1)
-    self.assertArraysEqual(f_out1, g_out1)
-
-    with self.assertRaisesRegex(
-        ValueError,
-        r"Compiled object called with input sharding.*does not match the "
-        r"sharding.*the computation was compiled with"):
-      g(x, y2)
-
   def test_dce_no_array(self):
     mesh = jtu.create_global_mesh((2,), ('x',))
     arr = jax.device_put(np.arange(8.), NamedSharding(mesh, P('x')))