diff --git a/docs/modules/distributed.rst b/docs/modules/distributed.rst
index de8a98a6..13c33d07 100644
--- a/docs/modules/distributed.rst
+++ b/docs/modules/distributed.rst
@@ -2,10 +2,68 @@
  Distributed
 #############
 
+*******
+ graph
+*******
 
+.. automodule:: anemoi.models.distributed.graph
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+************
+ khop_edges
+************
+
+.. automodule:: anemoi.models.distributed.khop_edges
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+..
+   *************
+
+..
+   primitives
+
+..
+   *************
+
+..
+   .. automodule:: anemoi.models.distributed.primitives
+
+..
+   :members:
+
+..
+   :no-undoc-members:
+
+..
+   :show-inheritance:
+
+********
+ shapes
+********
+
+.. automodule:: anemoi.models.distributed.shapes
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+*************
+ transformer
+*************
+
+.. automodule:: anemoi.models.distributed.transformer
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
 
+*******
+ utils
+*******
 
-.. automodule:: anemoi.models.distributed
+.. automodule:: anemoi.models.distributed.utils
    :members:
    :no-undoc-members:
    :show-inheritance: