Make inspect_array_sharding inside shard_map work with check_vma=True | check_rep=True. Fixes #23936

PiperOrigin-RevId: 750987053

Author: yashk2810

jax/_src/core.py

@@ -1894,7 +1894,10 @@ def str_short_aval(shape, dtype, mesh, spec, vma,
 def get_vma(vma, mesh):
   if mesh.empty:
     return vma
+  axis_env_names = get_axis_env().axis_names()
   for i in vma:
+    if i in axis_env_names and i not in mesh._name_to_type:
+      continue
     if mesh._name_to_type[i] != AxisType.Manual:
       raise ValueError(
           "Axes mentioned in `vma` field of ShapedArray should"

jax/_src/debugging.py

@@ -459,6 +459,7 @@ def _inspect_sharding_lowering_rule(ctx: mlir.LoweringRuleContext, value, *,
       mesh = mesh_lib.Mesh(np.array(devices).reshape(am.axis_sizes),
                            am.axis_names)
   elif isinstance(axis_context, sharding_impls.SPMDAxisContext):
+    mesh = axis_context.mesh
     devices = axis_context.mesh._flat_devices_tuple
   else:
     raise NotImplementedError(type(axis_context))
@@ -470,7 +471,8 @@ def _hlo_sharding_callback(hlo_sharding: xc.HloSharding):
     if mesh.empty:
       return callback(
           sharding_impls._op_sharding_to_pos_sharding(hlo_sharding, devices))
-    pspec = parse_flatten_op_sharding(hlo_sharding, mesh)[0]
+    pspec = (P() if hlo_sharding.is_manual() else
+             parse_flatten_op_sharding(hlo_sharding, mesh)[0])
     return callback(NamedSharding(mesh, pspec))
 
   if len(devices) == 1:

tests/debugging_primitives_test.py

@@ -25,6 +25,7 @@
 from jax._src import debugging
 from jax._src import dispatch
 from jax._src import test_util as jtu
+from jax.sharding import PartitionSpec as P
 import jax.numpy as jnp
 import numpy as np
 
@@ -1120,6 +1121,28 @@ def test_visualize_pmap_sharding(self):
     """)
     self.assertEqual(output(), expected)
 
+  def test_visualize_sharding_shard_map(self):
+    mesh = jtu.create_mesh((2,), 'x')
+
+    def f():
+      a = jnp.zeros(1000)
+      debugging.visualize_array_sharding(a)
+      return a
+
+    with jtu.capture_stdout() as output:
+      f()  # doesn't crash
+
+    with jtu.capture_stdout() as output:
+      jax.jit(f, out_shardings=jax.NamedSharding(mesh, P('x')))()  # doesn't crash
+
+    with jtu.capture_stdout() as output:
+      jax.shard_map(f, mesh=mesh, in_specs=P(None), out_specs=P("x"))()  # doesn't crash
+
+    with jtu.capture_stdout() as output:
+      jax.shard_map(f, mesh=mesh, in_specs=P(None), out_specs=P("x"),
+                    check_vma=False)()  # doesn't crash
+
+
 class InspectShardingTest(jtu.JaxTestCase):
 
   def test_inspect_sharding_is_called_in_pjit(self):