diff --git a/checkpoint/CHANGELOG.md b/checkpoint/CHANGELOG.md
index a9a484f5..232a3bc3 100644
--- 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
index 6ebdf581..44038cf9 100644
--- 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
index 92b466b1..27d4eec1 100644
--- 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
index 3fdb0c6c..1c83cc18 100644
--- 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:
\ No newline at end of file
diff --git a/docs/api_reference/checkpoint.tree.rst b/docs/api_reference/checkpoint.tree.rst
index 0201dfb4..be8bbd89 100644
--- 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
\ No newline at end of file
diff --git a/docs/api_reference/checkpoint.utils.rst b/docs/api_reference/checkpoint.utils.rst
deleted file mode 100644
index 300c240b..00000000
--- a/docs/api_reference/checkpoint.utils.rst
+++ /dev/null
@@ -1,40 +0,0 @@
-General Utilities
-======================================
-
-.. currentmodule:: orbax.checkpoint.utils
-
-.. automodule:: orbax.checkpoint.utils
-
-Aggregate file
-------------------------
-.. autofunction:: leaf_is_placeholder
-.. autofunction:: leaf_placeholder
-.. autofunction:: name_from_leaf_placeholder
-
-Directories
-------------------------
-.. autofunction:: cleanup_tmp_directories
-.. autofunction:: get_tmp_directory
-.. autofunction:: create_tmp_directory
-.. autofunction:: get_save_directory
-.. autofunction:: is_gcs_path
-
-Checkpoint steps
-------------------------
-.. autofunction:: step_from_checkpoint_name
-.. autofunction:: checkpoint_steps_paths
-.. autofunction:: checkpoint_steps
-.. autofunction:: any_checkpoint_step
-.. autofunction:: tmp_checkpoints
-.. autofunction:: lockdir
-.. autofunction:: is_locked
-.. autofunction:: are_locked
-
-Sharding
-------------------------
-.. autofunction:: fully_replicated_host_local_array_to_global_array
-
-Misc.
-------------------------
-.. autofunction:: is_scalar
-.. autofunction:: record_saved_duration
\ No newline at end of file
diff --git a/docs/api_reference/export.export_manager_base.rst b/docs/api_reference/export.export_manager_base.rst
deleted file mode 100644
index 73f78599..00000000
--- a/docs/api_reference/export.export_manager_base.rst
+++ /dev/null
@@ -1,11 +0,0 @@
-ExportManagerBase
-======================================
-
-.. currentmodule:: orbax.export.export_manager_base
-
-.. automodule:: orbax.export.export_manager_base
-
-ExportManagerBase
-------------------------
-.. autoclass:: ExportManagerBase
-  :members: