Skip to content

Commit 9f70133

Browse files
ChromeHeartsOrbax Authors
ChromeHearts
authored and
Orbax Authors
committed
Update V1 Docs
PiperOrigin-RevId: 842353538
1 parent ecb86c8 commit 9f70133

File tree

7 files changed

+92
-15
lines changed

7 files changed

+92
-15
lines changed

‎checkpoint/orbax/checkpoint/experimental/v1/serialization.py‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,3 @@
2828
SerializationContext,
2929
DeserializationContext,
3030
)
31-
32-
from orbax.checkpoint.experimental.v1._src.serialization.array_leaf_handler import (
33-
ArrayMetadata,
34-
)
35-
36-
from orbax.checkpoint.experimental.v1._src.serialization.numpy_leaf_handler import (
37-
NumpyShapeDtype,
38-
NumpyMetadata,
39-
)

‎docs/api_reference/checkpoint.v1.handlers.rst‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,25 @@
11
Checkpointable Handlers
22
============================================================================
33

4-
.. toctree::
5-
:maxdepth: 3
4+
.. currentmodule:: orbax.checkpoint.experimental.v1.handlers
65

7-
checkpoint.v1.handlers.types
8-
checkpoint.v1.handlers.pytree_handler
9-
checkpoint.v1.handlers.proto_handler
6+
.. automodule:: orbax.checkpoint.experimental.v1.handlers
7+
:members:
8+
9+
Types
10+
------------------------------------------------------------
11+
.. autoclass:: CheckpointableHandler
12+
.. autoclass:: StatefulCheckpointable
13+
14+
Handlers
15+
------------------------------------------------------------
16+
.. autoclass:: PyTreeHandler
17+
.. autoclass:: ProtoHandler
18+
.. autoclass:: JsonHandler
19+
20+
Registration
21+
------------------------------------------------------------
22+
.. autoclass:: CheckpointableHandlerRegistry
23+
.. autofunction:: global_registry
24+
.. autofunction:: local_registry
25+
.. autofunction:: register_handler

‎docs/api_reference/checkpoint.v1.partial.rst‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
Checkpointing Partial
2+
============================================================================
3+
4+
.. currentmodule:: orbax.checkpoint.experimental.v1.partial
5+
6+
.. automodule:: orbax.checkpoint.experimental.v1.partial
7+
:members:
8+
9+
Saving
10+
------------------------------------------------------------
11+
.. autofunction:: save_pytree
12+
.. autofunction:: save_pytree_async
13+
.. autofunction:: finalize

‎docs/api_reference/checkpoint.v1.path.rst‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
Checkpointing Path
2+
============================================================================
3+
4+
.. currentmodule:: orbax.checkpoint.experimental.v1.path
5+
6+
.. automodule:: orbax.checkpoint.experimental.v1.path
7+
:members:
8+
9+
Types
10+
------------------------------------------------------------
11+
.. autoclass:: Path
12+
:members:
13+
.. autoclass:: PathLike
14+
:members:
15+
.. autoclass:: PathAwaitingCreation
16+
:members:

‎docs/api_reference/checkpoint.v1.rst‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,7 @@ orbax.checkpoint API Reference
1212
checkpoint.v1.handler_registration
1313
checkpoint.v1.context
1414
checkpoint.v1.options
15-
15+
checkpoint.v1.serialization
16+
checkpoint.v1.path
17+
checkpoint.v1.tree
18+
checkpoint.v1.partial

‎docs/api_reference/checkpoint.v1.serialization.rst‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
Serialization
2+
============================================================================
3+
4+
.. currentmodule:: orbax.checkpoint.experimental.v1.serialization
5+
6+
.. automodule:: orbax.checkpoint.experimental.v1.serialization
7+
:members:
8+
9+
Registry
10+
------------------------------------------------------------
11+
.. autoclass:: BaseLeafHandlerRegistry
12+
:members:
13+
.. autoclass:: StandardLeafHandlerRegistry
14+
:members:
15+
16+
Types
17+
------------------------------------------------------------
18+
.. autoclass:: LeafHandler
19+
:members:
20+
.. autoclass:: SerializationParam
21+
:members:
22+
.. autoclass:: DeserializationParam
23+
:members:
24+
.. autoclass:: SerializationContext
25+
:members:
26+
.. autoclass:: DeserializationContext
27+
:members:

‎docs/api_reference/checkpoint.v1.tree.rst‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
Checkpointing Tree
2+
============================================================================
3+
4+
.. currentmodule:: orbax.checkpoint.experimental.v1.tree
5+
6+
.. automodule:: orbax.checkpoint.experimental.v1.tree
7+
:members:
8+
9+
Structure Utils
10+
------------------------------------------------------------
11+
.. autofunction:: merge

0 commit comments

Comments
 (0)