checkpoint
checkpoint
¶
Orbax's CheckpointManager is tad goofed. We ungoof it.
CheckpointManager()
¶
Save and restore full checkpoint paths with Orbax.
close() -> None
¶
Wait for outstanding saves and close the underlying checkpointer.
save(tree: PyTree[Any], path: str | os.PathLike[str]) -> None
¶
Start saving a pytree to a complete checkpoint path.
Waits for any previous save to finish before starting the new asynchronous save. All JAX processes must call this method collectively.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tree
|
PyTree[Any]
|
Pytree to checkpoint. |
required |
path
|
str | PathLike[str]
|
Complete checkpoint path. |
required |
restore(path: str | os.PathLike[str], template: PyTree[Any], allow_partial: bool = True) -> tuple[PyTree[Any], PyTree[bool] | None]
¶
Restore or surgically merge a pytree from a checkpoint.
All JAX processes must call this method collectively.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
path
|
str | PathLike[str]
|
Complete checkpoint path. |
required |
template
|
PyTree[Any]
|
Destination pytree defining structure, shapes, dtypes, and shardings. |
required |
allow_partial
|
bool
|
Whether destination leaves absent from, or incompatible with, the checkpoint may be initialized by the caller. |
True
|
Returns:
| Type | Description |
|---|---|
PyTree[Any]
|
The restored pytree and a destination-shaped Boolean pytree marking |
PyTree[bool] | None
|
leaves that require initialization. The mask is |
tuple[PyTree[Any], PyTree[bool] | None]
|
destination leaf was restored. |