code: def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):
      self.dense1 = nnx.Linear(din, dmid, rngs=rngs)
      self.dense2 = nnx.Linear(dmid, dout, rngs=rngs)
  
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 to create a new MLP object from inherited base classes: nnx with example: obj = MLP(din, dmid, dout)