Update TBE stable API docstrings

pytorch · Sep 25, 2024 · fd2524a · fd2524a
1 parent dad29b2
commit fd2524a
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