Update TBE stable API docstrings (#3171)

Summary: X-link: facebookresearch/FBGEMM#265 Update TBE docstrings and only show stable TBE APIs on the document page Pull Request resolved: #3171 Test Plan: See example https://deploy-preview-3171--pytorch-fbgemm-docs.netlify.app/fbgemm_gpu-python-api/table_batched_embedding_ops Reviewed By: shintaro-iwasaki Differential Revision: D63382237 Pulled By: sryap fbshipit-source-id: 287c2e4d9205faa16453b1ff6021433e0e4babde
pytorch · Sep 26, 2024 · 8e4bd7a · 8e4bd7a
1 parent f656ea8
commit 8e4bd7a
Show file tree

Hide file tree

Showing 3 changed files with 262 additions and 99 deletions.
diff --git a/fbgemm_gpu/docs/src/fbgemm_gpu-python-api/jagged_tensor_ops.rst b/fbgemm_gpu/docs/src/fbgemm_gpu-python-api/jagged_tensor_ops.rst
@@ -1,6 +1,8 @@
 Jagged Tensor Operators
 =======================
 
+.. automodule:: fbgemm_gpu
+
 .. autofunction:: torch.ops.fbgemm.jagged_2d_to_dense
 
 .. autofunction:: torch.ops.fbgemm.jagged_1d_to_dense

diff --git a/fbgemm_gpu/docs/src/fbgemm_gpu-python-api/table_batched_embedding_ops.rst b/fbgemm_gpu/docs/src/fbgemm_gpu-python-api/table_batched_embedding_ops.rst
@@ -1,5 +1,10 @@
-Table Batched Embedding (TBE) Operators
-=======================================
+Table Batched Embedding (TBE) Training Module
+=============================================
 
 .. autoclass:: fbgemm_gpu.split_table_batched_embeddings_ops_training.SplitTableBatchedEmbeddingBagsCodegen
-    :members:
+    :members: forward,
+              split_embedding_weights,
+              split_optimizer_states,
+              set_learning_rate,
+              update_hyper_parameters,
+              set_optimizer_step