Feature/Docs request: Saving and Loading checkpoints with optimizers and metrics

Dear Flax team,

I've been setting up all the infrastructure to save and load my flax models using Orbax and the guide in the [documentation](https://flax.readthedocs.io/en/latest/guides/checkpointing.html).

However, I noticed that a typical pattern when dealing with a flax model is to bundle the optimizer and the metrics:

```python3
tx = optax.chain(
...
)

graphdef, graphstate = nnx.split(
(model, nnx.Optimizer(model, tx, wrt=nnx.Param), metrics, ...)
)

model, optimizer, *rest = nnx.merge(graphdef, graphstate)
```

Using Orbrax is trivial to save and restore the graphstate. However, the graphdef is a whole other story. To reconstruct it, one would need to know exactly how it was created. What I've been doing so far is saving the import math of the model class and the constructor parameters as metadata to reconstruct the model, then splitting it, loading the graphstate, and replacing it. But when I was going to use it for a real task, I ran into the problem of reconstructing the graphdef when there is more than the model in there. The docs don't show how to do this.

The minor issue is separating the model graphstate from the rest of the things, and at least saving the model.

The larger issue would be how to manage the bundle of information in graphstate and graphdef.

A nice solution could be to create a to_json and from_json constructor for graphdef to make it easy to resialize with Orbax. This way, one would save the graphdef as metadata, and reconstruction would be as simple as loading both elements, constructing a new graphdef from the JSON, and then merging.

Best regards.