SPMD for initializing model using nnx.jit

[mmorinag127]

I would like to ask an example for initializing and updating a model using nnx.jit with SPMD.
Is there any relevant example?

[cgarciae]

nnx.jit's API is changing very soon (#3963) according to #4107. Currently you could use nnx.iter_graph to traverse the entire model and use with jax.device_put to shard the arrays inside Variables.

with path, parent in nnx.iter_graph(model):
  if isinstance(parent, nnx.Module):
    for name, value in vars(parent).items():
      if isinstance(value, nnx.Variable):
        sharding = ...
        value.value = jax.device_put(value.value, sharding)