Add stable API references (#2399)

Summary: Pull Request resolved: #2399 Differential Revision: D62700154
pytorch · Sep 17, 2024 · 8e5f839 · 8e5f839
1 parent 4fb2d00
commit 8e5f839
Show file tree

Hide file tree

Showing 9 changed files with 50 additions and 2 deletions.
diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -1,4 +1,5 @@
 sphinx==5.0.0
+myst_parser
 pyre-extensions
 # torch
 # PyTorch Theme

diff --git a/docs/source/_templates/layout.html b/docs/source/_templates/layout.html
@@ -4,7 +4,7 @@
 {{ super() }}
 
 <script type="text/javascript">
-  var collapsedSections = ['Introduction', 'All API References']
+  var collapsedSections = ['Introduction', 'Stable API References', 'All API References']
 </script>
 
 {% endblock %}
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -49,7 +49,11 @@
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
-extensions = ["sphinx.ext.napoleon", "sphinx.ext.autodoc"]
+extensions = [
+    "sphinx.ext.napoleon",
+    "sphinx.ext.autodoc",
+    "myst_parser",
+]
 
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -35,6 +35,17 @@ TorchRec API
 
    overview.rst
 
+.. toctree::
+   :maxdepth: 2
+   :caption: Stable API References
+   :hidden:
+
+   stable/datatypes-api-reference.rst
+   stable/modules-api-reference.rst
+   stable/planner-api-reference.rst
+   stable/model-parallel-api-reference.rst
+   stable/inference-api-reference.rst
+
 .. toctree::
    :maxdepth: 1
    :caption: All API References

diff --git a/docs/source/stable/datatypes-api-reference.rst b/docs/source/stable/datatypes-api-reference.rst
@@ -0,0 +1,10 @@
+TorchRec DataTypes API Reference
+----------------------------------
+
+.. automodule:: torchrec.sparse.jagged_tensor
+
+.. autoclass:: JaggedTensor
+
+.. autoclass:: KeyedJaggedTensor
+
+.. autoclass:: KeyedTensor
diff --git a/docs/source/stable/inference-api-reference.rst b/docs/source/stable/inference-api-reference.rst
@@ -0,0 +1,7 @@
+TorchRec Inference API Reference
+----------------------------------
+
+.. automodule:: torchrec.inference.modules
+
+.. autofunction:: quantize_inference_model
+.. autofunction:: shard_quant_model
diff --git a/docs/source/stable/model-parallel-api-reference.rst b/docs/source/stable/model-parallel-api-reference.rst
@@ -0,0 +1,6 @@
+TorchRec Model Parallel API Reference
+----------------------------------
+
+.. automodule:: torchrec.distributed.model_parallel
+
+.. autoclass:: DistributedModelParallel
diff --git a/docs/source/stable/modules-api-reference.rst b/docs/source/stable/modules-api-reference.rst
@@ -0,0 +1,7 @@
+TorchRec Modules API Reference
+----------------------------------
+
+.. automodule:: torchrec.modules.embedding_modules
+
+.. autoclass:: EmbeddingBagCollection
+.. autoclass:: EmbeddingCollection
diff --git a/docs/source/stable/planner-api-reference.rst b/docs/source/stable/planner-api-reference.rst
@@ -0,0 +1,2 @@
+TorchRec Planner API Reference
+----------------------------------