Modified documentation.

Author: vbharadwaj-bk

docs/api.rst

@@ -31,9 +31,26 @@ trying our code. OpenEquivariance cannot accelerate all tensor products; see
     :members:
     :undoc-members:
 
-.. autofunction:: openequivariance.impl_torch_to_oeq_dtype
+.. autofunction:: openequivariance.torch_to_oeq_dtype
+
+.. autofunction:: openequivariance.torch_ext_so_path
+
+OpenEquivariance JAX API
+------------------------
+The JAX API consists of ``TensorProduct`` and ``TensorProductConv``
+classes that behave identically to their PyTorch counterparts. These classes
+do not conform exactly to the e3nn-jax API, but perform the same computation. 
+
+.. autoclass:: openequivariance.jax.TensorProduct
+    :members: forward, reorder_weights_from_e3nn, reorder_weights_to_e3nn
+    :undoc-members:
+    :exclude-members:
+
+.. autoclass:: openequivariance.jax.TensorProductConv
+    :members: forward, reorder_weights_from_e3nn, reorder_weights_to_e3nn
+    :undoc-members:
+    :exclude-members: 
 
-.. autofunction:: openequivariance.impl_torch_ext_so_path
 
 API Identical to e3nn
 ---------------------

docs/conf.py

@@ -32,11 +32,14 @@
     "sphinx.ext.autodoc",
 ]
 
-sys.path.insert(0, str(Path("..").resolve()))
+sys.path.insert(0, str(Path("../openequivariance").resolve()))
 
 autodoc_mock_imports = [
     "torch",
+    "jax",
     "openequivariance.impl_torch.extlib",
+    "openequivariance.impl_jax.extlib",
+    "openequivariance_extjax",
     "jinja2",
     "numpy",
 ]

openequivariance/openequivariance/__init__.py

@@ -1,6 +1,6 @@
 # ruff: noqa: F401
 import sys
-import torch
+import os
 import numpy as np
 
 from pathlib import Path
@@ -13,12 +13,6 @@
     _MulIr,
     Instruction,
 )
-from openequivariance.impl_torch.TensorProduct import TensorProduct
-from openequivariance.impl_torch.TensorProductConv import (
-    TensorProductConv,
-)
-from openequivariance.impl_torch.extlib import torch_ext_so_path
-from openequivariance.core.utils import torch_to_oeq_dtype
 
 __version__ = None
 try:
@@ -44,20 +38,36 @@ def extension_source_path():
     """
     return str(Path(__file__).parent / "extension")
 
+TensorProduct, TensorProductConv, torch_ext_so_path, torch_to_oeq_dtype = None, None, None, None
 
-torch.serialization.add_safe_globals(
-    [
-        TensorProduct,
-        TensorProductConv,
-        TPProblem,
-        Irrep,
-        Irreps,
-        _MulIr,
-        Instruction,
-        np.float32,
-        np.float64,
-    ]
-)
+if "OEQ_NOTORCH" not in os.environ or os.environ["OEQ_NOTORCH"] != "1":
+    import torch
+    from openequivariance.impl_torch.TensorProduct import TensorProduct 
+    from openequivariance.impl_torch.TensorProductConv import TensorProductConv
+
+    from openequivariance.impl_torch.extlib import torch_ext_so_path
+    from openequivariance.core.utils import torch_to_oeq_dtype
+
+    torch.serialization.add_safe_globals(
+        [
+            TensorProduct,
+            TensorProductConv,
+            TPProblem,
+            Irrep,
+            Irreps,
+            _MulIr,
+            Instruction,
+            np.float32,
+            np.float64,
+        ]
+    )
+
+jax = None
+try:
+    import openequivariance_extjax
+    import openequivariance.impl_jax as jax
+except ImportError:
+    pass
 
 __all__ = [
     "TPProblem",
@@ -67,4 +77,5 @@ def extension_source_path():
     "torch_to_oeq_dtype",
     "_check_package_editable",
     "torch_ext_so_path",
+    "jax"
 ]

openequivariance/openequivariance/core/TensorProductBase.py

@@ -44,7 +44,7 @@ def reorder_weights_from_e3nn(self, weights, has_batch_dim: bool = True):
         Reorders weights from ``e3nn`` canonical order to the order used by ``oeq``.
 
         :param weights: Weights in ``e3nn`` canonical order, either an
-                        np.ndarray or a torch.Tensor. Tensor of dimensions ``[B, problem.weight_numel]``
+                        np.ndarray, torch.Tensor or JAX array. Tensor of dimensions ``[B, problem.weight_numel]``
                         when ``has_batch_dim=True``, otherwise of dimensions ``[problem.weight_numel]``.
 
         :param has_batch_dim: If ``True``, treats the first dimension of weights as a batch dimension. Default: ``True``.
@@ -57,8 +57,8 @@ def reorder_weights_to_e3nn(self, weights, has_batch_dim: bool = True):
         r"""
         Reorders weights from ``oeq`` canonical order to the order used by ``e3nn``.
 
-        :param weights: Weights in ``oeq`` canonical order, either an
-                        np.ndarray or a torch.Tensor. Tensor of dimensions ``[B, problem.weight_numel]``
+        :param weights: Weights in ``oeq`` canonical order, either a
+                        np.ndarray, torch.Tensor or JAX array. Tensor of dimensions ``[B, problem.weight_numel]``
                         when ``has_batch_dim=True``, otherwise of dimensions ``[problem.weight_numel]``.
 
         :param has_batch_dim: If ``True``, treats the first dimension of wieghts as a batch dimension. Default: ``True``.

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -65,9 +65,9 @@ class TensorProduct(LoopUnrollTP):
     :param problem: Specification of the tensor product.
     """
 
-    def __init__(self, config: TPProblem):
+    def __init__(self, problem: TPProblem):
         dp = extlib.DeviceProp(0)
-        super().__init__(config, dp, extlib.postprocess_kernel, torch_op=False)
+        super().__init__(problem, dp, extlib.postprocess_kernel, torch_op=False)
 
         self.attrs = {
             "kernel": self.jit_kernel,
@@ -78,7 +78,7 @@ def __init__(self, config: TPProblem):
         }
         hash_attributes(self.attrs)
 
-        self.weight_numel = config.weight_numel
+        self.weight_numel = problem.weight_numel
         self.L3_dim = self.config.irreps_out.dim
 
     def forward(self, X: jax.numpy.ndarray, Y: jax.numpy.ndarray, W: jax.numpy.ndarray) -> jax.numpy.ndarray: