What does in_specs=None mean in shard_map, and what is the difference with in_specs=P()?

[HeavyCrab]

In the document of jax.shard_map, the description of in_specs says (shard_map doc)

If None, inputs will be treated as static.

It is unclear what does "static" mean here.

According to my understanding, None should mean that the corresponding array will not be sharded. If it’s a Python data type such as a string, that’s easier to understand, since it only needs to keep one copy on the host side. But if it’s a jax.Array, what’s the difference between this and P()? If I mark it as None, does that imply some GPU-to-GPU communication occurs at every operation, making it less efficient than P()?

From my experiments, this does indeed seem to be the case. See the code below:

import jax
import jax.numpy as jnp
from functools import partial
from jax.sharding import PartitionSpec as P
import timeit

mesh = jax.make_mesh((2,), ('d',))

@partial(jax.shard_map, mesh=mesh, in_specs=(None, P()), out_specs=P())
def f(x, y):
    z = 0
    for i in range(10):
        z = z + x + y
    return z

@partial(jax.shard_map, mesh=mesh, in_specs=(P(), P()), out_specs=P())
def g(x, y):
    z = 0
    for i in range(10):
        z = z + x + y
    return z


x = jnp.array(1)
y = jnp.array(2)
z1 = f(x, y)
z2 = g(x, y)

print("start")

print("f")
print(timeit.timeit(lambda: jax.block_until_ready(f(x, y)), number=50))
print("g")
print(timeit.timeit(lambda: jax.block_until_ready(g(x, y)), number=50))

The output is:

start
f
60.49180791340768
g
49.207496081944555

[mattjj]

Thanks for the question!

I agree we shouldn't use the term "static" here.

An in_spec of None means to treat that input the same as if it had been closed over, i.e.

@jax.shard_map(in_specs=(None, <A>), <B>)
def f(x, <C>):
  <D>

f(x, <E>)

means the same thing as

@jax.shard_map(in_specs=<A>, <B>)
def f_(<C>):
  <D>

f_(<E>)

Everything else follows from that, but just to be concrete:

But if it’s a jax.Array, what’s the difference between this and P()?

There's not much of a difference, because closed-over arraylikes are treated as if they have specs of P().

If I mark it as None, does that imply some GPU-to-GPU communication occurs at every operation, making it less efficient than P()?

Under a jit, there should be no difference. Indeed putting a jax.jit on your two functions makes the timing identical for me.

Outside of a jit, i.e. in eager mode as in your example, if the value in question is not on device (e.g. if it's a numpy.ndarray) passing something as an input with spec P() will ensure it's transferred to device only once. If the value is already an Array on device, as in your example, we might still be doing some reshaping operations multiple times for f that we only do once for g. But I'd consider that a missing optimization that we could implement if it becomes important.

What do you think?

[HeavyCrab]

Thank you so much for your reply :) That makes sense. Some clarification in the document would make it perfect.
p.s.: the doc for jax.experimental.shard_map.shard_map (doc) is a bit different with jax.shard_map. In that document a None in_specs is explained as

If an argument, or argument subtree, has a corresponding spec of None, that argument is not sharded

If these two APIs are identical, maybe it is better to update.