Access full shape of input arguments inside shard_map

[Findus23]

While looking through the new shard_map examples I had a thought about a potential feature that I think is missing. I am not sure how often it would be useful, but I think it could make debugging more clear.

Let's assume you have a multidimensional array that is sharded along one axis and you pass it to a function using shard_map to interact with the data along the sharded axis.
Then the array x inside shard_map will have a smaller shape along the axis that is sharded. But when writing more complex code inside shard_map I think it might be harder to keep track of which arrays are "normal" arrays and which ones are "shard_map views into sharded arrays". x.vma does tell you that x is shard_mapped, but it does not give us the information which of the dimensions of x is. And I guess this will get worse if there are multiple sharded dimensions.

# using jax 0.6.1.dev20250428
import os

import jax
from jax import shard_map
from jax.experimental.shard import reshard
from jax.sharding import PartitionSpec as P, AxisType, set_mesh

os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={8}'

mesh = jax.make_mesh((4,), ("gpus",), axis_types=(AxisType.Explicit,))
set_mesh(mesh)

input_sharded = reshard(jax.numpy.zeros((8, 8, 8)), P(None, "gpus", None))

@shard_map(
    mesh=mesh,
    in_specs=P(None, 'gpus', None),
    out_specs=P(None, 'gpus', None)
)
def f(x):
    print(jax.typeof(x))  # float32[8,2,8]{gpus}
    print(x.vma)  # frozenset({'gpus'})
    print(jax.typeof(x).sharding.spec)  # PartitionSpec(None, None, None)
    print(????(x)) # -> index: 1 (or {"gpus": 1})
    return x

f(input_sharded)

Or in other words: If this is easily possible and others find this useful to, would it be possible to also provide the information that gpus applies along the second axis of x to x?

[yashk2810]

Well inside, shard_map, nothing is sharded anymore so all you can do is track which mesh axes you are varying across.

Can you explain more about your use case? Maybe there is a different way to do what you want.

[Findus23]

That makes sense.

I don't really have a more specific use case and was just wondering if I was missing a simple way to know which of the axes is the varying one. At least for me personally it helps quite a bit keeping track of which array has which properties during programming to not get confused in the calculations (see also https://github.com/Findus23/jax-array-info/).

But it's not really important as indeed in shard_map the arrays are considered unsharded.

[yashk2810]

I guess the best way to do what you want, is to track shapes then :)

Or you can enter into shard_map one manual axis at a time. That will make sure other axes are explicit/auto, thus giving you information about sharding for those axes. Here's an example:

def test_partial_auto_explicit(self, mesh):
    mesh = jax.make_mesh((2, 2), ('i', 'j'), axis_types=(AxisType.Explicit, AxisType.Explicit))
 
    def g(x):
      self.assertDictEqual(x.aval.sharding.mesh._axis_types_dict,
                           {AxisType.Manual: ('i',), AxisType.Explicit: ('j',)})
      # You can see how `x` is sharded on `j` mesh axes since that axis is Explicit!
      self.assertEqual(x.aval.sharding.spec, P(None, 'j'))
      out = x * x
      self.assertEqual(out.aval.sharding.spec, P(None, 'j'))
      return out

    @jax.jit
    def f(x):
      x = jax.shard_map(g, out_specs=P('i', None), axis_names={'i'})(x)
      self.assertEqual(x.aval.sharding.spec, P('i', 'j'))
      return x

    v = jnp.arange(32.).reshape(4, 8)
    v = jax.device_put(v, jax.sharding.NamedSharding(mesh, P('i', 'j')))
    
    with jax.sharding.use_mesh(mesh):
      out = f(v)