V1

• Introduction to Checkpointing with Orbax
  • Reading and Writing
  • Checkpointing in a Training Loop
  • API Overview
  • What’s Next?
• Orbax v0 to v1 Migration Guide
  • How to use v1 API to load checkpoints saved with v0 API
    • Checkpoint Layouts
    • Compatibility Matrix
  • Migrating from v0 CheckpointManager to v1 Checkpointer
• Working with Checkpointables
  • Key Concepts
    • Checkpoints as a Bundle
    • The “Checkpointable” Concept
    • Checkpointable vs. Abstract Checkpointable
  • Concrete Examples
    • Saving and Loading Multiple Checkpointables
    • Interoperating with save / load
    • Partial Loading
    • Understanding Abstract Checkpointables
    • Working with Real Datasets
• Working with PyTree Checkpoints
  • Exclusive APIs to checkpoint PyTrees
  • Basic Checkpointing
    • Understanding Abstract Trees and Leaves
    • Standard Leaf Types
    • Customizing Loaded Properties for Arrays
    • Partial Loading
  • Advanced Customizations
    • Saving
    • Loading
• Checkpoint Format Guide
  • What is an Orbax checkpoint?
  • Overview
    • Singular Checkpoints
    • Sequence of Checkpoints
  • Format Details
    • Generic Checkpoints
    • PyTree Checkpointables
    • Other Checkpointables
• Checkpointing in a Training Loop
  • Getting Started
  • Checkpointer APIs
    • Saving
    • Querying Available Checkpoints
    • Inspecting Checkpoint Metadata
    • Garbage Collection
    • Loading
    • Checkpointables and Dataset Checkpointing
  • Training with MNIST Data
• Asynchronous Checkpointing
  • What is this?
  • Why should I care?
  • Usage
  • Additional Details
• Customizing Checkpointing Behavior
  • Custom Checkpointables
    • Implementing Point as a StatefulCheckpointable
    • Supporting Point with CheckpointableHandler
• Partial Saving
  • The Core Concept: The Partial Save Session
  • API and Basic Usage: Adding to a Checkpoint
  • Code Example: A Simple Addition Workflow
    • Initial Save
    • Add More Data
    • Aside: Loading Before Finalizing
    • Finalize the Checkpoint
    • Verify the Result
  • API Reference
  • Advanced Workflow: Combining Partial Saving and Partial Restore
    • Use Case: Creating an Inference-Ready Checkpoint
    • Setup
    • The Efficient Update Workflow
    • Update and Partial Save
    • Verification
  • Atomicity Guarantees
• Model Surgery
  • Partial Loading
    • Placeholder
    • Omission
    • Model Surgery
• Checkpointing and Exporting JAX Models: An End-to-End Guide with Orbax
  • 1. Setup
    • Installation
    • Imports
    • Helper for Directory Management
  • 2. Define a Simulated JAX State
    • Define a Model and Training State
  • 3. Orbax Checkpointing Workflow
    • Create a Checkpointing Directory
    • Checkpointing During a Simulated Training Loop
    • Resuming from a Checkpoint
  • 4. Saving Final JAX Parameters for Export
    • Extract Final Parameters for Saving
    • Using save_pytree for the Final Save
    • Loading Exported Parameters (Verification)
  • 5. Exporting to TensorFlow SavedModel
    • Define JAX Model Apply Function and Pre/Post-processing for Export
    • Create JaxModule and ServingConfig
    • Export to TensorFlow SavedModel
    • Validate the Exported Model
• orbax.checkpoint V1 API Reference
  • Submodules
    • ocp.v1.arrays module
    • ocp.v1.errors module
    • ocp.v1.handlers module
    • CheckpointableHandler Registration
    • ocp.v1.options module
    • ocp.v1.partial module
    • ocp.v1.path module
    • ocp.v1.training module
    • ocp.v1.tree module
    • ocp.v1.multihost module
  • Top-level Symbols
    • Types
    • Loading
    • Saving
    • Metadata
    • Path Utilities
    • Synchronization
    • Checkpointable Handlers
    • Context