Add stable API references (#2399)

Summary:
Pull Request resolved: #2399

Remove old API references that was very noisy. Insert API references for specific stable APIs.

Reviewed By: iamzainhuda

Differential Revision: D62700154

fbshipit-source-id: 231a69decdf5c347ee895bc548a237a87982b167

Author: PaulZhang12

docs/source/_templates/layout.html

@@ -4,7 +4,7 @@
 {{ super() }}
 
 <script type="text/javascript">
-  var collapsedSections = ['Introduction', 'Getting Started', 'All API References']
+  var collapsedSections = ['Introduction', 'Getting Started', 'API References']
 </script>
 
 {% endblock %}

docs/source/datatypes-api-reference.rst

@@ -0,0 +1,22 @@
+Data Types
+-------------------
+
+
+TorchRec contains data types for representing embedding, otherwise known as sparse features.
+Sparse features are typically indices that are meant to be fed into embedding tables. For a given
+batch, the number of embedding lookup indices are variable. Therefore, there is a need for a **jagged**
+dimension to represent the variable amount of embedding lookup indices for a batch.
+
+This section covers the classes for the 3 TorchRec data types for representing sparse features:
+**JaggedTensor**, **KeyedJaggedTensor**, and **KeyedTensor**.
+
+.. automodule:: torchrec.sparse.jagged_tensor
+
+.. autoclass:: JaggedTensor
+    :members:
+
+.. autoclass:: KeyedJaggedTensor
+    :members:
+
+.. autoclass:: KeyedTensor
+    :members:

docs/source/index.rst

@@ -100,20 +100,12 @@ how you can contribute:
    setup-torchrec.rst
 
 .. toctree::
-   :maxdepth: 1
-   :caption: All API References
+   :maxdepth: 2
+   :caption: API References
    :hidden:
 
-   torchrec.datasets.rst
-   torchrec.datasets.scripts.rst
-   torchrec.distributed.rst
-   torchrec.distributed.planner.rst
-   torchrec.distributed.sharding.rst
-   torchrec.fx.rst
-   torchrec.inference.rst
-   torchrec.models.rst
-   torchrec.modules.rst
-   torchrec.optim.rst
-   torchrec.quant.rst
-   torchrec.sparse.rst
-   torchrec.metrics.rst
+   datatypes-api-reference.rst
+   modules-api-reference.rst
+   planner-api-reference.rst
+   model-parallel-api-reference.rst
+   inference-api-reference.rst

docs/source/inference-api-reference.rst

@@ -0,0 +1,19 @@
+Inference
+----------------------------------
+
+TorchRec provides easy-to-use APIs for transforming an authored TorchRec model
+into an optimized inference model for distributed inference, via eager module swaps.
+
+This transforms TorchRec modules like ``EmbeddingBagCollection`` in the model to
+a quantized, sharded version that can be compiled using torch.fx and TorchScript
+for inference in a C++ environment.
+
+The intended use is calling ``quantize_inference_model`` on the model followed by
+``shard_quant_model``.
+
+.. codeblock::
+
+.. automodule:: torchrec.inference.modules
+
+.. autofunction:: quantize_inference_model
+.. autofunction:: shard_quant_model

docs/source/model-parallel-api-reference.rst

@@ -0,0 +1,10 @@
+Model Parallel
+----------------------------------
+
+``DistributedModelParallel`` is the main API for distributed training with TorchRec optimizations.
+
+
+.. automodule:: torchrec.distributed.model_parallel
+
+.. autoclass:: DistributedModelParallel
+    :members:

docs/source/modules-api-reference.rst

@@ -0,0 +1,30 @@
+Modules
+----------------------------------
+
+Standard TorchRec modules represent collections of embedding tables:
+
+* ``EmbeddingBagCollection`` is a collection of ``torch.nn.EmbeddingBag``
+* ``EmbeddingCollection`` is a collection of ``torch.nn.Embedding``
+
+These modules are constructed through standardized config classes:
+
+* ``EmbeddingBagConfig`` for ``EmbeddingBagCollection``
+* ``EmbeddingConfig`` for ``EmbeddingCollection``
+
+.. automodule:: torchrec.modules.embedding_configs
+
+.. autoclass:: EmbeddingBagConfig
+    :show-inheritance:
+
+.. autoclass:: EmbeddingConfig
+    :show-inheritance:
+
+.. autoclass:: BaseEmbeddingConfig
+
+.. automodule:: torchrec.modules.embedding_modules
+
+.. autoclass:: EmbeddingBagCollection
+    :members:
+
+.. autoclass:: EmbeddingCollection
+    :members:

docs/source/planner-api-reference.rst

@@ -0,0 +1,50 @@
+Planner
+----------------------------------
+
+The TorchRec Planner is responsible for determining the most performant, balanced
+sharding plan for distributed training and inference.
+
+The main API for generating a sharding plan is ``EmbeddingShardingPlanner.plan``
+
+.. automodule:: torchrec.distributed.types
+
+.. autoclass:: ShardingPlan
+    :members:
+
+.. automodule:: torchrec.distributed.planner.planners
+
+.. autoclass:: EmbeddingShardingPlanner
+    :members:
+
+.. automodule:: torchrec.distributed.planner.enumerators
+
+.. autoclass:: EmbeddingEnumerator
+    :members:
+
+.. automodule:: torchrec.distributed.planner.partitioners
+
+.. autoclass:: GreedyPerfPartitioner
+    :members:
+
+
+.. automodule:: torchrec.distributed.planner.storage_reservations
+
+.. autoclass:: HeuristicalStorageReservation
+    :members:
+
+.. automodule:: torchrec.distributed.planner.proposers
+
+.. autoclass:: GreedyProposer
+    :members:
+
+
+.. automodule:: torchrec.distributed.planner.shard_estimators
+
+.. autoclass:: EmbeddingPerfEstimator
+    :members:
+
+
+.. automodule:: torchrec.distributed.planner.shard_estimators
+
+.. autoclass:: EmbeddingStorageEstimator
+    :members: