add equinox modules to docs

e3nn · Jan 23, 2025 · be9056f · be9056f
1 parent e81640d
commit be9056f
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 1 deletion.
diff --git a/docs/api/equinox.rst b/docs/api/equinox.rst
@@ -0,0 +1,5 @@
+Equinox Modules
+===============
+
+.. autoclass:: e3nn_jax.equinox.Linear
+    :members:
diff --git a/docs/api/index.rst b/docs/api/index.rst
@@ -13,6 +13,7 @@ API
     radial
     haiku
     flax
+    equinox
     s2
     extra
     utils

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -13,3 +13,4 @@ jraph
 nequip_jax @ git+https://github.com/mariogeiger/[email protected]
 flax
 dm-haiku
+equinox
diff --git a/e3nn_jax/_src/linear_equinox.py b/e3nn_jax/_src/linear_equinox.py
@@ -28,7 +28,7 @@ def _get_gradient_normalization(
 
 
 class Linear(eqx.Module):
-    r"""Equivariant Linear Flax module
+    r"""Equivariant Linear Equinox module
 
     Args:
         irreps_out (`Irreps`): output representations, if allowed bu Schur's lemma.
-Original file line number
+Diff line change
@@ Expand Up / @@ -13,6 +13,7 @@ API @@
         radial
         haiku
         flax
+        equinox
         s2
         extra
         utils
@@ Expand Down @@