Merge pull request #97 from teddykoker/eqx-docs

ameya98 · web-flow · commit 403f0ef159c5 · 2025-01-23T16:45:38.000-05:00
Add equinox modules to docs
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/nequip-jax@1.1.0
 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
     radial
     haiku
     flax
 +    equinox
     s2
     extra
     utils