Multi-host migration of `nnx.pmap` to `nnx.shard_map`

### Context

**TL;DR:** Is there a way to use `nnx.shard_map` to replicate `nnx.pmap` with host-local dataloading and sharding?

First, thanks @cgarciae for adding `nnx.shard_map`! I'm trying to migrate from `nnx.pmap` to `nnx.shard_map`. Here's the gist of my current `pmap` code:

```python
# def step_fn(train_state, data): ...

PMAP_AXIS = "pmap_axis"
state_axes = nnx.StateAxes({MyCustomState: 0, nnx.RngState: 0, ..., None})
pmapped_fn = nnx.split_rngs(jax.local_device_count())(
nnx.pmap(step_fn, axis_name=PMAP_AXIS, in_axes=(state_axes, 0))
)
```

I have a custom state that I manually create outside of `pmap` of shape `(num_local_devices, ...)`. I also use per-host dataloading, where the data is of shape `(num_local_devices, ...)`. The `0`s in the state axes informs `pmap` of these leading local device axes, which are sharded accordingly.

### Issue

With multi-host training, `pmap` is convenient because:

1. `lax` collectives aggregate across hosts per `jax.pmap` [specs](https://docs.jax.dev/en/latest/_autosummary/jax.pmap.html);
2. I can therefore have per-host logic (dataloading, custom state) while still having inter-host operations.

I'm struggling to see how I can leverage `nnx.shard_map` to achieve the same. Here is my attempt:

```python
# def step_fn(train_state, data): ...

# setup mesh
HOST_AXIS = "host_axis"
DEVICE_AXIS = "device_axis"
num_hosts = jax.process_count()
num_local_devices = jax.local_device_count()
device_array = np.asarray(jax.devices()).reshape((num_hosts, num_local_devices))
mesh = jax.sharding.Mesh(device_array, (HOST_AXIS, DEVICE_AXIS))

# sharding
SHARD = P((HOST_AXIS, DEVICE_AXIS))
REPLICATE = P()
state_sharding = nnx.StateSharding({MyCustomState: SHARD, ...: REPLICATE})
sharded_fn = nnx.shard_map(step_fn, mesh=mesh, in_specs=(state_sharding, SHARD), out_specs=(REPLICATE, REPLICATE))
# optional: jitted_fn = nnx.jit(sharded_fn)
```

There are a few issues:

1. I lose `pmap`'s "per-host code with inter-host collective" property. I need the `mesh` to be a global mesh in order to properly leverage `lax` collectives. This means that I need to use functions like JAX's [`host_local_array_to_global_array`](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.multihost_utils.host_local_array_to_global_array.html) or [`make_array_from_process_local_data`](https://docs.jax.dev/en/latest/_autosummary/jax.make_array_from_process_local_data.html) to build a global array view, then shard across the global mesh. These "globalization" ops also seem to incur non-trivial overhead especially if I want to write to and update some per-host state.
2. This makes it difficult to have any per-host logic. For instance, if I want to have some custom per-host behavior in my dataloader or `MyCustomState`, this becomes very difficult due to the need to work with a global mesh.
3. (minor) `nnx.split_rngs` does not work nicely with `shard_map` due to differences in dimensions.

### Desired Solultion

To make things as clear as possible, here's what I think I want to do:

```python
# host local
local_devices_array = np.array(jax.local_devices())
local_mesh = jax.sharding.Mesh(local_devices_array, (DEVICE_AXIS,))
HOST_LOCAL_SHARD = NamedSharding(local_mesh, P(DEVICE_AXIS))

# host global
num_hosts = jax.process_count()
num_local_devices = jax.local_device_count()
device_array = np.asarray(jax.devices()).reshape((num_hosts, num_local_devices))
global_mesh = jax.sharding.Mesh(device_array, (HOST_AXIS, DEVICE_AXIS))
GLOBAL_REPLICATE = NamedSharding(global_mesh, P(HOST_AXIS, DEVICE_AXIS))

# shard_map (want; mix and match host local and global semantics)
state_sharding = nnx.StateSharding({MyCustomState: HOST_LOCAL_SHARD, ...: GLOBAL_REPLICATE})
sharded_fn = nnx.shard_map(step_fn, mesh=mesh, in_specs=(state_sharding, HOST_LOCAL_SHARD), out_specs=(GLOBAL_REPLICATE, GLOBAL_REPLICATE))
```

To the best of my knowledge, `jax.pmap` actually uses `jax.shard_map` under the hood ([source](https://github.com/jax-ml/jax/blob/main/jax/_src/api.py#L1485)), which makes me believe this transition should be possible.

Thanks in advance!