Batch test is working.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -90,6 +90,7 @@ def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
         return reorder_jax(self.forward_schedule, weights, "backward", not self.config.shared_weights)
 
     def forward_cpu(self, L1_in, L2_in, L3_out, weights) -> None:
+        weights = self.reorder_weights_from_e3nn(weights, has_batch_dim=not self.config.shared_weights)
         result = self.forward(
             jax.numpy.asarray(L1_in),
             jax.numpy.asarray(L2_in),
@@ -100,6 +101,7 @@ def forward_cpu(self, L1_in, L2_in, L3_out, weights) -> None:
     def backward_cpu(
         self, L1_in, L1_grad, L2_in, L2_grad, L3_grad, weights, weights_grad
     ) -> None:
+        weights = self.reorder_weights_from_e3nn(weights, has_batch_dim=not self.config.shared_weights)
         backward_fn = jax.vjp(
             lambda X, Y, W: self.forward(X, Y, W),
             jax.numpy.asarray(L1_in),
@@ -111,4 +113,5 @@ def backward_cpu(
         )
         L1_grad[:] = np.asarray(L1_grad_jax)
         L2_grad[:] = np.asarray(L2_grad_jax)
-        weights_grad[:] = np.asarray(weights_grad_jax)
+        weights_grad[:] = np.asarray(weights_grad_jax)
+        weights_grad[:] = self.reorder_weights_to_e3nn(weights_grad, has_batch_dim=not self.config.shared_weights)

tests/batch_test.py

@@ -254,7 +254,9 @@ def problem(self, request, dtype):
 
 class TestTorchbindDisable(TestProductionModels):
     @pytest.fixture(scope="class")
-    def extra_tp_constructor_args(self):
+    def extra_tp_constructor_args(self, test_jax):
+        if test_jax:
+            pytest.skip("N/A for JAX")
         return {"use_opaque": True}
 
 
@@ -268,11 +270,14 @@ def problem(self, request, dtype):
         return problem
 
     @pytest.fixture(scope="class")
-    def tp_and_problem(self, problem, extra_tp_constructor_args):
-        tp = TensorProduct(problem, **extra_tp_constructor_args)
-        switch_map = {
-            np.float32: torch.float64,
-            np.float64: torch.float32,
-        }
-        tp.to(switch_map[problem.irrep_dtype])
-        return tp, tp.config
+    def tp_and_problem(self, problem, extra_tp_constructor_args, test_jax):
+        if test_jax:
+            pytest.skip("N/A for JAX")
+        else:
+            tp = oeq.TensorProduct(problem, **extra_tp_constructor_args)
+            switch_map = {
+                np.float32: torch.float64,
+                np.float64: torch.float32,
+            }
+            tp.to(switch_map[problem.irrep_dtype])
+            return tp, tp.config