Backward convolution is failing, need to figure out why.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -107,5 +107,3 @@ def backward_cpu(
         L2_grad[:] = np.asarray(L2_grad_jax)
         weights_grad[:] = np.asarray(weights_grad_jax)
 
-
-

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -162,3 +162,52 @@ def __call__(
         sender_perm: Optional[jax.numpy.ndarray] = None,
     ) -> jax.numpy.ndarray:
         return self.forward(X, Y, W, rows, cols, sender_perm)
+
+    def forward_cpu(self, L1_in, L2_in, weights, L3_out, graph):
+        rows = graph.rows.astype(np.int32)
+        cols = graph.cols.astype(np.int32)
+        sender_perm = graph.transpose_perm.astype(np.int32)
+        result = self.forward(
+            jax.numpy.asarray(L1_in),
+            jax.numpy.asarray(L2_in),
+            jax.numpy.asarray(weights),
+            jax.numpy.asarray(rows),
+            jax.numpy.asarray(cols),
+            jax.numpy.asarray(sender_perm),
+        )
+        L3_out[:] = np.asarray(result)
+
+    def backward_cpu(
+        self,
+        L1_in,
+        L1_grad,
+        L2_in,
+        L2_grad,
+        L3_grad,
+        weights,
+        weights_grad,
+        graph,
+    ):
+        rows = graph.rows.astype(np.int32)
+        cols = graph.cols.astype(np.int32)
+        sender_perm = graph.transpose_perm.astype(np.int32)
+
+        backward_fn = jax.vjp(
+            lambda X, Y, W: self.forward(
+                X,
+                Y,
+                W,
+                jax.numpy.asarray(rows),
+                jax.numpy.asarray(cols),
+                jax.numpy.asarray(sender_perm),
+            ),
+            jax.numpy.asarray(L1_in),
+            jax.numpy.asarray(L2_in),
+            jax.numpy.asarray(weights),
+        )[1]
+        L1_grad_jax, L2_grad_jax, weights_grad_jax = backward_fn(
+            jax.numpy.asarray(L3_grad)
+        )
+        L1_grad[:] = np.asarray(L1_grad_jax)
+        L2_grad[:] = np.asarray(L2_grad_jax)
+        weights_grad[:] = np.asarray(weights_grad_jax)

tests/conftest.py

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

tests/conv_test.py

@@ -4,6 +4,7 @@
 from pytest_check import check
 
 import numpy as np
+import openequivariance
 import openequivariance as oeq
 from openequivariance.benchmark.ConvBenchmarkSuite import load_graph
 from itertools import product
@@ -51,23 +52,32 @@ def graph(self, request):
     def extra_conv_constructor_args(self):
         return {}
 
+    @pytest.fixture(scope="class")
+    def test_jax(self, request):
+        return request.config.getoption("--jax")
+
     @pytest.fixture(params=["atomic", "deterministic", "kahan"], scope="class")
-    def conv_object(self, request, problem, extra_conv_constructor_args):
+    def conv_object(self, request, problem, extra_conv_constructor_args, test_jax):
+        cls = oeq.TensorProductConv
+        if test_jax:
+            from openequivariance.impl_jax import TensorProductConv as jax_conv
+            cls = jax_conv
+
         if request.param == "atomic":
-            return oeq.TensorProductConv(
+            return cls(
                 problem, deterministic=False, **extra_conv_constructor_args
             )
         elif request.param == "deterministic":
             if not problem.shared_weights:
-                return oeq.TensorProductConv(
+                return cls(
                     problem, deterministic=True, **extra_conv_constructor_args
                 )
             else:
                 pytest.skip("Shared weights not supported with deterministic")
         elif request.param == "kahan":
             if problem.irrep_dtype == np.float32:
                 if not problem.shared_weights:
-                    return oeq.TensorProductConv(
+                    return cls(
                         problem,
                         deterministic=True,
                         kahan=True,