Can shard_map be nested?

[yoeldr]

Suppose I have a function thats calls shard_map, is it possible to parallaize it using shard_map (or by other means) over some unused axes?

For example, consider shard_loss_by_model below, which shards the computation of loss over the 'model' axis. Is it possible to parallelize it over the 'batch' axis?

import chex
import jax
from jax.experimental import mesh_utils
from jax.experimental.shard_map import shard_map
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P


chex.set_n_cpu_devices(8)
devices = mesh_utils.create_device_mesh((4, 2))
mesh = Mesh(devices, axis_names=('batch', 'model'))

w = jax.device_put(
    jnp.arange(2, dtype=jnp.float32),
    NamedSharding(mesh, P('model')),
)
x = jax.device_put(
    jnp.arange(16, dtype=jnp.float32),
    NamedSharding(mesh, P('batch')),
)


def loss(w, x):
  result = jnp.sum((x.reshape(-1, 1) - w) ** 2)
  return jax.lax.psum(result, axis_name='model')


shard_loss_by_model = shard_map(
    loss, mesh=mesh, in_specs=(P('model'), P()), out_specs=P()
)

Note that the naive implementation (see nested below) throws an AttributeError: 'RewriteTrace' object has no attribute 'process_shard_map' exception.

nested = shard_map(
    shard_loss_by_model,
    mesh=mesh,
    in_specs=(P(), P('batch')),
    out_specs=P('batch'),
)
print(nested(w, x))

[jeffgortmaker]

Any chance you've figured this out since posting? Running into the same error in a similar setting.