Fix Orbax readthedoc build

PiperOrigin-RevId: 694570187
google · Nov 8, 2024 · 4ddd6ba · 4ddd6ba
1 parent d8a2a63
commit 4ddd6ba
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 80 deletions.
diff --git a/checkpoint/CHANGELOG.md b/checkpoint/CHANGELOG.md
@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.9.0] - 2024-11-08
+
 ### Changed
 - Create `Composite` class, which `CompositeArgs` now subclasses.
 - Move `tree` to `_src`.
@@ -16,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  actionable info.
 - Make `NameFormat.find_all` impls concurrent.
 - Move `path` package under `_src` package.
+- Updated readthedoc
 
 ### Fixed
 - Emergency checkpoint: use JAX for global_max and combine multiple broadcasts

diff --git a/checkpoint/orbax/checkpoint/version.py b/checkpoint/orbax/checkpoint/version.py
@@ -16,7 +16,7 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
-__version__ = '0.8.0'
+__version__ = '0.9.0'
 
 
 # TODO: b/362813406 - Add latest change timestamp and commit number.

diff --git a/docs/api_reference/checkpoint.logging.rst b/docs/api_reference/checkpoint.logging.rst
@@ -8,30 +8,30 @@ Checkpoint logging
 
 AbstractLogger
 ------------------------
-.. autoclass:: AbstractLogger
+.. autoclass:: abstract_logger.AbstractLogger
   :members:
 
 CompositeLogger
 ------------------------
-.. autoclass:: CompositeLogger
+.. autoclass:: composite_logger.CompositeLogger
   :members:
 
 StandardLogger
 ------------------------
-.. autoclass:: StandardLogger
+.. autoclass:: standard_logger.StandardLogger
   :members:
 
 CloudLogger
 ------------------------
-.. autoclass:: CloudLogger
+.. autoclass:: cloud_logger.CloudLogger
   :members:
 
 CloudLoggerOptions
 ------------------------
-.. autoclass:: CloudLoggerOptions
+.. autoclass:: cloud_logger.CloudLoggerOptions
   :members:
 
 StepStatistics
 ------------------------
-.. autoclass:: StepStatistics
+.. autoclass:: step_statistics.StepStatistics
   :members:
diff --git a/docs/api_reference/checkpoint.metadata.rst b/docs/api_reference/checkpoint.metadata.rst
@@ -9,9 +9,6 @@ Metadata Utilities
 Tree Metadata
 ------------------------
 
-.. autoclass:: Metadata
-  :members:
-
 .. autoclass:: ArrayMetadata
   :members:
 
@@ -21,6 +18,10 @@ Tree Metadata
 .. autoclass:: StringMetadata
   :members:
 
+.. autoclass:: StorageMetadata
+  :members:
+
+
 Sharding Metadata
 ------------------------
 
@@ -37,11 +38,4 @@ Sharding Metadata
   :members:
 
 .. autoclass:: PositionalShardingMetadata
-  :members:
-
-.. autoclass:: ShardingTypes
-  :members:
-
-.. autofunction:: from_jax_sharding
-.. autofunction:: from_serialized_string
-.. autofunction:: get_sharding_or_none
+  :members:
diff --git a/docs/api_reference/checkpoint.tree.rst b/docs/api_reference/checkpoint.tree.rst
@@ -5,16 +5,9 @@ Tree Utilities
 
 .. automodule:: orbax.checkpoint.tree
 
-.. autofunction:: deserialize_tree
-.. autofunction:: from_flat_dict
-.. autofunction:: from_flattened_with_keypath
+.. autofunction:: get_param_names
 .. autofunction:: serialize_tree
+.. autofunction:: deserialize_tree
 .. autofunction:: to_flat_dict
-.. autofunction:: is_sequence_key
-.. autofunction:: is_dict_key
-.. autofunction:: tuple_path_from_keypath
-.. autofunction:: get_key_name
-.. autofunction:: is_empty_node
-.. autofunction:: is_empty_or_leaf
-.. autofunction:: to_shape_dtype_struct
-.. autofunction:: get_param_names
+.. autofunction:: from_flat_dict
+.. autofunction:: to_shape_dtype_struct
diff --git a/docs/api_reference/checkpoint.utils.rst b/docs/api_reference/checkpoint.utils.rst
diff --git a/docs/api_reference/export.export_manager_base.rst b/docs/api_reference/export.export_manager_base.rst