code: def create_model(seed: int):
      return MLP(10, 20, 30, rngs=nnx.Rngs(seed))
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/08_save_load_checkpoints.py
imports: from tempfile import TemporaryDirectory
  import jax
  import jax.numpy
  import orbax.checkpoint
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling a global function with example usage: create_model(seed) and returns: MLP