Remove jax_varying_axes_in_types config and rewrite from shard_map_p

PiperOrigin-RevId: 748545142

Author: yashk2810

jax/_src/config.py

@@ -235,7 +235,6 @@ def trace_context():
           threefry_partitionable.value,
           threefry_gpu_kernel_lowering.value,
           use_direct_linearize.value,
-          varying_axes_in_types.value,
           softmax_custom_jvp.value,
           disable_jit.value,
           debug_key_reuse.value,
@@ -1092,14 +1091,6 @@ def enum_flag(name, default, *args, **kwargs) -> Flag[str]:
     help=('Use direct linearization instead JVP followed by partial eval'),
     include_in_jit_key=True)
 
-varying_axes_in_types = bool_state(
-    name='jax_varying_axes_in_types',
-    default=True,
-    help=('Adds varying manual axes to ShapedArray to track which mesh axes the'
-          ' array is varying over. This will help to remove the efficient'
-          ' transpose rewrite machinery in shard_map'),
-    include_in_jit_key=True)
-
 # TODO make it so people don't use this, this is internal...
 _check_rep = bool_state(
     name='check_rep',

jax/_src/core.py

@@ -2006,8 +2006,6 @@ def pvary(x, axis_name):
 pvary_p.def_impl(lambda *args, axes, axis_index_groups: args)
 
 def _pvary_abstract_eval(*args, axes, axis_index_groups):
-  if not config.varying_axes_in_types.value:
-    return args
   if not config._check_rep.value:
     return args
   assert isinstance(axes, tuple)
@@ -2027,8 +2025,6 @@ def _pvary_abstract_eval(*args, axes, axis_index_groups):
 
 
 def standard_insert_pvary(*args):
-  if not config.varying_axes_in_types.value:
-    return args
   if not config._check_rep.value:
     return args
   if not args:
@@ -2040,8 +2036,6 @@ def standard_insert_pvary(*args):
           if out_vma - src else arg for arg, src in zip(args, in_vma)]
 
 def standard_vma_rule(prim_name, *avals, **kwargs) -> frozenset[AxisName]:
-  if not config.varying_axes_in_types.value:
-    return frozenset()
   if not config._check_rep.value:
     return frozenset()
   avals = tuple(a for a in avals if a is not abstract_token)

jax/_src/interpreters/batching.py

@@ -407,7 +407,7 @@ def __init__(self, trace, val, batch_dim: NotMapped | int | RaggedAxis,
   def aval(self):
     aval = core.get_aval(self.val)
     if self._trace.axis_data.spmd_name is not None:
-      if config._check_rep.value and config.varying_axes_in_types.value:
+      if config._check_rep.value:
         aval = aval.update(
             vma=aval.vma - frozenset(self._trace.axis_data.spmd_name))
     if self.batch_dim is not_mapped:
@@ -776,7 +776,7 @@ def _batch_jaxpr2(
       aval = core.unmapped_aval(
           axis_data.size, b, aval, axis_data.explicit_mesh_axis)
       if axis_data.spmd_name is not None:
-        if config._check_rep.value and config.varying_axes_in_types.value:
+        if config._check_rep.value:
           aval = aval.update(vma=aval.vma | frozenset(axis_data.spmd_name))  # type: ignore
       avals_in2.append(aval)
   jaxpr_out, _, consts, () = pe.trace_to_jaxpr_dynamic(f, avals_in2)
@@ -1111,7 +1111,7 @@ def broadcast(x, sz, axis, mesh_axis=None):
   # out how to ensure jaxpr arguments always have the context mesh.
   with mesh_lib.use_abstract_mesh(sharding.mesh):
     x = jax.lax.broadcast_in_dim(x, shape, broadcast_dims, out_sharding=sharding)
-    if config._check_rep.value and config.varying_axes_in_types.value:
+    if config._check_rep.value:
       # TODO(yashkatariya,parkers): don't do this, fix during fixit week 2026
       spmd_names = core.get_axis_env().spmd_axis_names
       if len(spmd_names) > 1:

jax/_src/lax/parallel.py

@@ -144,7 +144,7 @@ def pos_reduce(x):
       size = math.prod([core.get_axis_env().axis_size(name) for name in named_axes])
     out_flat = tuple(lax._const(leaf, size) * pos_reduce(leaf) for leaf in leaves)
   else:
-    if config.varying_axes_in_types.value and config._check_rep.value:
+    if config._check_rep.value:
       out_flat = bind_psum_invariant(
           leaves, axes=tuple(axis_name), axis_index_groups=axis_index_groups)
     else:
@@ -827,9 +827,6 @@ def _allreduce_effectful_abstract_eval(*args, axes, axis_index_groups):
   return out_avals, {core.NamedAxisEffect(axis) for axis in named_axes}
 
 def _psum_invariant_abstract_eval(name, *args, axes, axis_index_groups):
-  if not config.varying_axes_in_types.value:
-    return psum_p.abstract_eval(
-        *args, axes=axes, axis_index_groups=axis_index_groups)
   if not config._check_rep.value:
     return psum_p.abstract_eval(
         *args, axes=axes, axis_index_groups=axis_index_groups)
@@ -865,9 +862,6 @@ def _psum_invariant_abstract_eval(name, *args, axes, axis_index_groups):
 
 # TODO(yashkatariya): Replace this with _psum_invariant_abstract_eval
 def _pmin_pmax_abstract_eval(name, *args, axes, axis_index_groups):
-  if not config.varying_axes_in_types.value:
-    return _allreduce_effectful_abstract_eval(
-        *args, axes=axes, axis_index_groups=axis_index_groups)
   if not config._check_rep.value:
     return _allreduce_effectful_abstract_eval(
         *args, axes=axes, axis_index_groups=axis_index_groups)
@@ -1417,8 +1411,6 @@ def _ragged_all_to_all_batched_collective(axis_data, vals_in, dims_in,
 batching.skippable_batchers[ragged_all_to_all_p] = partial(_names_in_param, 'axis_name')
 
 def insert_collective_pvary(axis_name, x):
-  if not config.varying_axes_in_types.value:
-    return x
   if not config._check_rep.value:
     return x
 
@@ -1551,8 +1543,6 @@ def _all_gather_lowering(ctx, x, *, all_gather_dimension, axis_name,
 
 
 def collective_vma_rule(prim_name, axis_name, x_aval):
-  if not config.varying_axes_in_types.value:
-    return frozenset()
   if not config._check_rep.value:
     return frozenset()
   axis_name = (axis_name,) if not isinstance(axis_name, tuple) else axis_name
@@ -1921,8 +1911,7 @@ def _axis_index_effectful_abstract_eval(*, axis_name):
   mesh = get_abstract_mesh()
   sharding = NamedSharding(mesh, P())
   vma = ((frozenset(axis_name) if mesh._any_axis_manual else frozenset())
-         if config.varying_axes_in_types.value and config._check_rep.value
-         else frozenset())
+         if config._check_rep.value else frozenset())
   return ShapedArray((), np.int32, sharding=sharding, vma=vma), effect
 
 def _axis_index_batcher(axis_data, vals_in, dims_in, *, axis_name):