p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_lam.py (python)
    model_state = nnx.state(optimizer.model)
    model_sharded_state = jax.device_put(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.device_put(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)
[content] train_dynamics.py (python)
    model_state = nnx.state(optimizer.model)
    model_sharded_state = jax.device_put(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.device_put(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)