code: def load_model(path: str) -> MLP:
      # create that model with abstract shapes
      state, static = jax.eval_shape(lambda: create_model(0).split())
      # Load the parameters
      checkpointer = orbax.PyTreeCheckpointer()
      state = checkpointer.restore(f"{path}/state", item=state)
      # Merge the parameters into the model
      model = static.merge(state)
      return model
  
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 3 for calling a global function with example usage: load_model(path) and returns: model