Deprecation Warnings with orbax 0.5.3

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Debian 12
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: flax 0.8.1, orbax-checkpoint 0.5.3
- Python version: 3.11.4

### Problem you have encountered:
I followed [save and load checkpoints tutorial](https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html) and I got deprecation warnings. Although checkpoints were saved correctly, it would be great if the latest correct way of saving/loading Flax `TraningState` is documented in the tutorial.

### What you expected to happen:
no warnings

### Logs, error messages, etc:

```
WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/api_refactor.html to migrate by May 1st, 2024.
WARNING:absl:SaveArgs.aggregate is deprecated, please use custom TypeHandler (https://orbax.readthedocs.io/en/latest/custom_handlers.html#typehandler) or contact Orbax team to migrate before May 1st, 2024.
```

### Steps to reproduce:
Just follow [save and load checkpoints tutorial](https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html)