Skip to main content
Ctrl+K

Orbax documentation

Checkpointing

  • Announcements
  • Checkpointing with Orbax
  • API Overview
  • Using the Refactored CheckpointManager API
  • Checkpointing PyTrees of Arrays
  • Checkpoint Format Guide
  • Optimized Checkpointing with Tensorstore
  • Transformations
  • Preemption Tolerance
  • Asynchronous Checkpointing
  • Debugging Guide
  • orbax.checkpoint API Reference
    • AbstractCheckpointManager
    • CheckpointManager
    • Checkpointers
    • CheckpointHandlers
    • CheckpointArgs
    • TypeHandlers
    • Atomicity
    • Atomicity Types
    • Step entities
    • Checkpoint Deleter
    • Preservation Policy
    • Save Decision Policy
    • Checkpointing Utilities
    • PyTree Transformations
    • Msgpack utils
    • Multi–host Utilities
    • Tree Utilities
    • Metadata Utilities
    • Checkpoint logging
    • Checkpointing Options
    • Handler Registration
    • Arrays
  • V1
    • Introduction to Checkpointing with Orbax
    • Orbax v0 to v1 Migration Guide
    • Working with Checkpointables
    • Working with PyTree Checkpoints
    • Checkpoint Format Guide
    • Checkpointing in a Training Loop
    • Asynchronous Checkpointing
    • Customizing Checkpointing Behavior
    • Partial Saving
    • Model Surgery
    • Checkpointing and Exporting JAX Models: An End-to-End Guide with Orbax
    • orbax.checkpoint V1 API Reference
      • 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
        • Training Checkpointer
        • ocp.v1.training.save_decision_policies module
        • ocp.v1.training.errors module
        • Training Types
      • ocp.v1.tree module
      • ocp.v1.multihost module

Exporting

  • Exporting with Orbax
  • orbax.export API Reference
    • DTensor utilities for multi–device/host export
    • ExportManager
    • JaxModule
    • ServingConfig
    • General Utilities
    • ValidationManager
    • ValidationReport

Additional Information

  • Contributors
  • .md

Contributors

Contents

  • Core Contributors
    • Checkpointing
    • Exporting
    • General
  • Community Guidelines

Contributors#

Orbax does not currently accept contributions from the open source community.

Core Contributors#

Checkpointing#

  • Colin Gaffney

Exporting#

  • Dinghua Li

  • John Zhang

General#

  • Ruoxin Sang

  • Ayush Jain

  • Haitang Hu

A full list of contributors can be found here.

Community Guidelines#

This project follows Google’s Open Source Community Guidelines.

previous

ValidationReport

Contents
  • Core Contributors
    • Checkpointing
    • Exporting
    • General
  • Community Guidelines

By Orbax Contributors

© Copyright 2026, Google.