docs(distributed): add distributed API reference

XuehaiPan · XuehaiPan · commit ffdaa0f2db6a · 2023-02-15T14:32:59.000Z
diff --git a/docs/source/api/api.rst b/docs/source/api/api.rst
@@ -285,6 +285,115 @@ Chain
 .. autofunction:: chain
 
 
+Distributed Utilities
+=====================
+
+.. currentmodule:: torchopt.distributed
+
+Initialization and Synchronization
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autosummary::
+
+    auto_init_rpc
+    barrier
+
+.. autofunction:: auto_init_rpc
+.. autofunction:: barrier
+
+Process group information
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autosummary::
+
+    get_world_info
+    get_world_rank
+    get_rank
+    get_world_size
+    get_local_rank
+    get_local_world_size
+    get_worker_id
+
+.. autofunction:: get_world_info
+.. autofunction:: get_world_rank
+.. autofunction:: get_rank
+.. autofunction:: get_world_size
+.. autofunction:: get_local_rank
+.. autofunction:: get_local_world_size
+.. autofunction:: get_worker_id
+
+Worker selection
+~~~~~~~~~~~~~~~~
+
+.. autosummary::
+
+    on_rank
+    not_on_rank
+    rank_zero_only
+    rank_non_zero_only
+
+.. autofunction:: on_rank
+.. autofunction:: not_on_rank
+.. autofunction:: rank_zero_only
+.. autofunction:: rank_non_zero_only
+
+Remote Procedure Call (RPC)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autosummary::
+
+    remote_async_call
+    remote_sync_call
+
+.. autofunction:: remote_async_call
+.. autofunction:: remote_sync_call
+
+Predefined partitioners and reducers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autosummary::
+
+    dim_partitioner
+    batch_partitioner
+    mean_reducer
+    sum_reducer
+
+.. autofunction:: dim_partitioner
+.. autofunction:: batch_partitioner
+.. autofunction:: mean_reducer
+.. autofunction:: sum_reducer
+
+Function parallelization wrappers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autosummary::
+
+    parallelize
+    parallelize_async
+    parallelize_sync
+
+.. autofunction:: parallelize
+.. autofunction:: parallelize_async
+.. autofunction:: parallelize_sync
+
+Distributed Autograd
+~~~~~~~~~~~~~~~~~~~~
+
+.. currentmodule:: torchopt.distributed
+
+.. autosummary::
+
+    context
+    get_gradients
+    backward
+    grad
+
+.. autofunction:: context
+.. autofunction:: get_gradients
+.. autofunction:: backward
+.. autofunction:: grad
+
+
 General Utilities
 =================
 
diff --git a/docs/source/distributed/distributed.rst b/docs/source/distributed/distributed.rst
@@ -142,7 +142,6 @@ Initialization and Synchronization
 
 .. autosummary::
 
-
     torchopt.distributed.auto_init_rpc
     torchopt.distributed.barrier
 
@@ -197,7 +196,6 @@ Process group information
 
 .. autosummary::
 
-
     torchopt.distributed.get_world_info
     torchopt.distributed.get_world_rank
     torchopt.distributed.get_rank
@@ -228,7 +226,6 @@ Worker selection
 
 .. autosummary::
 
-
     torchopt.distributed.on_rank
     torchopt.distributed.not_on_rank
     torchopt.distributed.rank_zero_only
@@ -275,7 +272,6 @@ Remote Procedure Call (RPC)
 
 .. autosummary::
 
-
     torchopt.distributed.remote_async_call
     torchopt.distributed.remote_sync_call
 
@@ -354,7 +350,6 @@ Predefined partitioners and reducers
 
 .. autosummary::
 
-
     torchopt.distributed.dim_partitioner
     torchopt.distributed.batch_partitioner
     torchopt.distributed.mean_reducer
@@ -439,7 +434,6 @@ Function parallelization wrappers
 
 .. autosummary::
 
-
     torchopt.distributed.parallelize
     torchopt.distributed.parallelize_async
     torchopt.distributed.parallelize_sync
@@ -490,7 +484,6 @@ Distributed Autograd
 
 .. autosummary::
 
-
     torchopt.distributed.autograd.context
     torchopt.distributed.autograd.get_gradients
     torchopt.distributed.autograd.backward
