Developed and delivered checkpointing functionality for Keras3 Jax TPU Embeddings in the keras-team/keras-rs repository, focusing on robust save and restore capabilities for embeddings, metrics, and training state. Leveraged Python, JAX, and Orbax to implement a custom CheckpointManager and a Keras training loop Callback, integrating checkpointing seamlessly into standard workflows. This work enabled reliable training resumption and improved state management for long-running TPU experiments, reducing downtime and supporting experiment continuity. The solution addressed the need for resilient model checkpointing in deep learning workflows, demonstrating depth in model state handling and practical integration with existing machine learning infrastructure and tools.