Replace auto tracking with manual axis names tracking internally in shard_map.

This will make it easier to implement nested shard_map where you enter into manual mode one axis at a time per shard_map.

PiperOrigin-RevId: 751509665

Author: yashk2810

jax/_src/checkify.py

@@ -968,15 +968,15 @@ def shard_map_error_check(
   new_in_names = (*([{}] * num_error_vals), *in_names)
   new_vals_in = [*err_vals, *vals_in]
   in_avals = list(map(core.get_aval, new_vals_in))
-  auto = kwargs.get('auto')
+  manual_axes = kwargs.get('manual_axes')
   check_vma = kwargs.get('check_vma')
   for i, v in enumerate(in_avals):
     if not (sharder := core.shard_aval_handlers.get(type(v))):
       raise ValueError(f'Unsupported aval type: {type(v)}')
-    in_avals[i] = sharder(mesh, auto, check_vma, new_in_names[i], v)
+    in_avals[i] = sharder(mesh, manual_axes, check_vma, new_in_names[i], v)
 
-  with (jshmap._extend_axis_env(mesh, auto),
-        mesh_lib.use_abstract_mesh(jshmap._as_manual_mesh(mesh, auto)),
+  with (jshmap._extend_axis_env(mesh, manual_axes),
+        mesh_lib.use_abstract_mesh(jshmap._as_manual_mesh(mesh, manual_axes)),
         config._check_vma(check_vma)):
     # jaxpr to checked_jaxpr
     checked_jaxpr, out_tree, _ = jaxpr_to_checkify_jaxpr(

jax/_src/interpreters/pxla.py

@@ -3005,8 +3005,6 @@ def from_hlo(name: str,
         allow_prop_to_outputs, tuple(host_callbacks), backend, da, pmap_nreps,
         compiler_options_kvs, pgle_profiler)
 
-    orig_out_shardings = out_shardings
-
     if auto_spmd_lowering:
       assert mesh is not None
       in_shardings_xla, out_shardings_xla = _get_mesh_pspec_shardings_from_executable(