Keras Tutorial: Checkpointing distributed models with Orbax

A walkthrough on integrating Orbax checkpointing with Keras models using the JAX backend, specifically for multi-host distributed training environments. Two utility classes are needed: a custom Orbax checkpoint manager that wraps Keras's internal get/set state tree APIs to save and restore model weights, optimizer state, and metrics, and an Orbax checkpoint callback that hooks into Keras's training loop to auto-save at epoch end and restore on training start. This approach fills a gap in Keras's built-in checkpointing, which does not support multi-host setups, and ensures no more than one epoch of work is lost on failure.