Add equinox modules to docs #97

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

ameya98 merged 1 commit into e3nn:main from teddykoker:eqx-docs

Jan 23, 2025

docs/api/equinox.rst

-Original file line number
+Diff line change
@@ -0,0 +1,5 @@
+    Equinox Modules
+    ===============
+    .. autoclass:: e3nn_jax.equinox.Linear
+        :members:

docs/api/index.rst

-Original file line number
+Diff line change
@@ Expand Up / @@ -13,6 +13,7 @@ API @@
         radial
         haiku
         flax
+        equinox
         s2
         extra
         utils
@@ Expand Down @@

docs/requirements.txt

-Original file line number
+Diff line change
@@ Expand Up / @@ -13,3 +13,4 @@ jraph @@
     nequip_jax @ git+https://github.com/mariogeiger/nequip-jax@1.1.0
     flax
     dm-haiku
+    equinox

e3nn_jax/_src/linear_equinox.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -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.
@@ Expand Down @@

Provide feedback