{"library":"orbax-checkpoint","title":"Orbax Checkpoint","description":"Orbax Checkpoint provides a robust, asynchronous, and fault-tolerant checkpointing library primarily designed for JAX and Flax models. It allows saving and restoring large-scale machine learning model states and arbitrary data structures efficiently, supporting distributed environments and custom serialization. It is currently at version 0.11.34 and frequently updates to align with JAX/Flax developments.","language":"python","status":"active","last_verified":"Wed May 13","install":{"commands":["pip install orbax-checkpoint jax flax","pip install orbax-checkpoint"],"cli":null},"imports":["import orbax.checkpoint as ocp","from orbax.checkpoint import CheckpointManager","from orbax.checkpoint import CheckpointManagerOptions","from orbax.checkpoint import StandardSave","from orbax.checkpoint import StandardRestore"],"auth":{"required":false,"env_vars":[]},"quickstart":{"code":"import jax\nimport jax.numpy as jnp\nimport orbax.checkpoint as ocp\nimport os\nimport shutil\n\n# Define a temporary checkpoint directory\nckpt_dir = '/tmp/my_orbax_checkpoint_example'\nif os.path.exists(ckpt_dir):\n    shutil.rmtree(ckpt_dir)\nos.makedirs(ckpt_dir, exist_ok=True)\n\n# 1. Create a CheckpointManager\noptions = ocp.CheckpointManagerOptions(\n    save_interval_steps=1,\n    max_to_keep=3,\n    keep_time_interval_secs=None\n)\nmngr = ocp.CheckpointManager(ckpt_dir, options=options)\n\n# 2. Prepare some data to save\nstep = 0\ndata_to_save = {'params': jnp.array([1.0, 2.0, 3.0])}\nprint(f\"Saving data at step {step}: {data_to_save['params']}\")\n\n# 3. Save the checkpoint\n# Ensure to wrap data with StandardSave for explicit serialization\nmngr.save(step, args=ocp.StandardSave(data_to_save))\nmngr.wait_until_finished() # Ensure save completes\n\n# Simulate more steps and saves\nstep = 1\ndata_to_save = {'params': jnp.array([4.0, 5.0, 6.0])}\nprint(f\"Saving data at step {step}: {data_to_save['params']}\")\nmngr.save(step, args=ocp.StandardSave(data_to_save))\nmngr.wait_until_finished()\n\n# 4. Restore the latest checkpoint\nlatest_step = mngr.latest_step()\nif latest_step is not None:\n    print(f\"\\nRestoring data from latest step: {latest_step}\")\n    # Provide a template for StandardRestore, even if just the expected structure\n    restored_data = mngr.restore(latest_step, args=ocp.StandardRestore(data_to_save))\n    print(f\"Restored data: {restored_data['params']}\")\nelse:\n    print(\"No checkpoint found to restore.\")\n\n# 5. Close the manager\nmngr.close()\n\n# Clean up\nif os.path.exists(ckpt_dir):\n    shutil.rmtree(ckpt_dir)\n","lang":"python","description":"Demonstrates how to initialize a CheckpointManager, save JAX array data, and restore the latest checkpoint. Highlights the use of `ocp.StandardSave` and `ocp.StandardRestore` for explicit serialization arguments and the importance of `wait_until_finished()`.","tag":null,"tag_description":null,"last_tested":"2026-04-24","results":[{"runtime":"python:3.10-alpine","exit_code":1},{"runtime":"python:3.10-slim","exit_code":1},{"runtime":"python:3.11-alpine","exit_code":1},{"runtime":"python:3.11-slim","exit_code":1},{"runtime":"python:3.12-alpine","exit_code":1},{"runtime":"python:3.12-slim","exit_code":1},{"runtime":"python:3.13-alpine","exit_code":1},{"runtime":"python:3.13-slim","exit_code":1},{"runtime":"python:3.9-alpine","exit_code":1},{"runtime":"python:3.9-slim","exit_code":1}]},"compatibility":{"tag":null,"tag_description":null,"last_tested":"2026-05-13","installed_version":"0.6.4","pypi_latest":"0.11.39","is_stale":true,"summary":{"python_range":"3.10–3.9","success_rate":63,"avg_install_s":15.5,"avg_import_s":3.47,"wheel_type":"wheel"},"results":[{"runtime":"python:3.10-alpine","python_version":"3.10","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"17.8M"},{"runtime":"python:3.10-alpine","python_version":"3.10","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.10-alpine","python_version":"3.10","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.10-alpine","python_version":"3.10","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.10-slim","python_version":"3.10","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":14.5,"import_time_s":1.73,"mem_mb":42.2,"disk_size":"680M"},{"runtime":"python:3.10-slim","python_version":"3.10","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":1.87,"mem_mb":42,"disk_size":"680M"},{"runtime":"python:3.10-slim","python_version":"3.10","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":16.3,"import_time_s":1.79,"mem_mb":42.2,"disk_size":"700M"},{"runtime":"python:3.10-slim","python_version":"3.10","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":1.77,"mem_mb":42,"disk_size":"700M"},{"runtime":"python:3.11-alpine","python_version":"3.11","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"19.6M"},{"runtime":"python:3.11-alpine","python_version":"3.11","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.11-alpine","python_version":"3.11","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.11-alpine","python_version":"3.11","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.11-slim","python_version":"3.11","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":14.9,"import_time_s":3.8,"mem_mb":53.4,"disk_size":"734M"},{"runtime":"python:3.11-slim","python_version":"3.11","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":4.27,"mem_mb":53.3,"disk_size":"734M"},{"runtime":"python:3.11-slim","python_version":"3.11","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":16.2,"import_time_s":3.82,"mem_mb":53.9,"disk_size":"758M"},{"runtime":"python:3.11-slim","python_version":"3.11","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":3.95,"mem_mb":53.7,"disk_size":"757M"},{"runtime":"python:3.12-alpine","python_version":"3.12","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"11.5M"},{"runtime":"python:3.12-alpine","python_version":"3.12","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.12-alpine","python_version":"3.12","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.12-alpine","python_version":"3.12","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.12-slim","python_version":"3.12","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":14.6,"import_time_s":4.47,"mem_mb":51.4,"disk_size":"721M"},{"runtime":"python:3.12-slim","python_version":"3.12","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":5.37,"mem_mb":51.2,"disk_size":"721M"},{"runtime":"python:3.12-slim","python_version":"3.12","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":15.8,"import_time_s":4.43,"mem_mb":52.7,"disk_size":"745M"},{"runtime":"python:3.12-slim","python_version":"3.12","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":5.42,"mem_mb":52.6,"disk_size":"744M"},{"runtime":"python:3.13-alpine","python_version":"3.13","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"11.2M"},{"runtime":"python:3.13-alpine","python_version":"3.13","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.13-alpine","python_version":"3.13","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.13-alpine","python_version":"3.13","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.13-slim","python_version":"3.13","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":14.3,"import_time_s":4.07,"mem_mb":52.8,"disk_size":"720M"},{"runtime":"python:3.13-slim","python_version":"3.13","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":5.61,"mem_mb":52.6,"disk_size":"719M"},{"runtime":"python:3.13-slim","python_version":"3.13","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":16.2,"import_time_s":4.22,"mem_mb":53.2,"disk_size":"743M"},{"runtime":"python:3.13-slim","python_version":"3.13","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":4.87,"mem_mb":53,"disk_size":"743M"},{"runtime":"python:3.9-alpine","python_version":"3.9","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"17.3M"},{"runtime":"python:3.9-alpine","python_version":"3.9","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.9-alpine","python_version":"3.9","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.9-alpine","python_version":"3.9","os_libc":"alpine (musl)","variant":"orbax-checkpoint","exit_code":1,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.9-slim","python_version":"3.9","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":15.1,"import_time_s":1.76,"mem_mb":42,"disk_size":"621M"},{"runtime":"python:3.9-slim","python_version":"3.9","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":2.29,"mem_mb":42,"disk_size":"621M"},{"runtime":"python:3.9-slim","python_version":"3.9","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":16.9,"import_time_s":1.69,"mem_mb":42,"disk_size":"641M"},{"runtime":"python:3.9-slim","python_version":"3.9","os_libc":"slim (glibc)","variant":"orbax-checkpoint","exit_code":0,"wheel_type":null,"failure_reason":null,"import_side_effects":null,"install_time_s":null,"import_time_s":2.12,"mem_mb":42,"disk_size":"641M"}]}}