1/3 tests is passing.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -1,4 +1,5 @@
 import jax
+import numpy as np
 from functools import partial
 from openequivariance.impl_jax import extlib
 from openequivariance.core.e3nn_lite import TPProblem
@@ -74,10 +75,24 @@ def __init__(self, config: TPProblem):
         self.weight_numel = config.weight_numel
         self.L3_dim = self.config.irreps_out.dim
 
-    def forward(self, X: jax.ndarray, Y: jax.ndarray, W: jax.ndarray) -> jax.ndarray:
+    def forward(self, X: jax.numpy.ndarray, Y: jax.numpy.ndarray, W: jax.numpy.ndarray) -> jax.numpy.ndarray:
         return forward(X, Y, W, self.L3_dim, self.config.irrep_dtype, self.attrs)
 
     def __call__(
         self, X: jax.numpy.ndarray, Y: jax.numpy.ndarray, W: jax.numpy.ndarray
     ) -> jax.numpy.ndarray:
         return self.forward(X, Y, W)
+    
+    def forward_cpu(
+        self,
+        L1_in: np.ndarray,
+        L2_in: np.ndarray,
+        L3_out: np.ndarray,
+        weights: np.ndarray,
+    ) -> None:
+        result = self.forward(
+            jax.numpy.asarray(L1_in),
+            jax.numpy.asarray(L2_in),
+            jax.numpy.asarray(weights),
+        )
+        L3_out[:] = np.asarray(result)

openequivariance/openequivariance/impl_jax/__init__.py

@@ -0,0 +1,4 @@
+from openequivariance.impl_jax.TensorProduct import TensorProduct as TensorProduct
+from openequivariance.impl_jax.TensorProductConv import TensorProductConv as TensorProductConv
+
+__all__ = ["TensorProduct", "TensorProductConv"]

tests/batch_test.py

@@ -2,8 +2,8 @@
 from pytest_check import check
 
 import numpy as np
+import openequivariance
 import openequivariance as oeq
-from openequivariance.impl_torch.TensorProduct import TensorProduct
 from openequivariance.benchmark.correctness_utils import (
     correctness_forward,
     correctness_backward,
@@ -19,7 +19,6 @@
 from itertools import product
 import torch
 
-
 class TPCorrectness:
     def thresh(self, direction):
         return {"fwd": 1e-5, "bwd": 3e-4, "double_bwd": 3e-4}[direction]
@@ -41,8 +40,16 @@ def extra_tp_constructor_args(self):
         return {}
 
     @pytest.fixture(scope="class")
-    def tp_and_problem(self, problem, extra_tp_constructor_args):
-        tp = TensorProduct(problem, **extra_tp_constructor_args)
+    def test_jax(self, request):
+        return request.config.getoption("--jax")
+
+    @pytest.fixture(scope="class")
+    def tp_and_problem(self, problem, extra_tp_constructor_args, test_jax):
+        cls = oeq.TensorProduct 
+        if test_jax:
+            import openequivariance.impl_jax.TensorProduct as jax_tp
+            cls = jax_tp
+        tp = cls(problem, **extra_tp_constructor_args)
         return tp, problem
 
     def test_tp_fwd(self, tp_and_problem):

tests/conftest.py

@@ -0,0 +1,8 @@
+import pytest
+import os
+
+os.environ["JAX_ENABLE_X64"] = "True"
+def pytest_addoption(parser):
+    parser.addoption(
+        "--jax", action="store", default=False, help="Test the JAX frontend instead of PyTorch"
+    )