Almost there, need to get TensorProductConv working.

Author: vbharadwaj-bk

openequivariance/openequivariance/core/ConvolutionBase.py

@@ -578,19 +578,19 @@ def test_correctness_double_backward(
             in1, in2, out_grad, weights, weights_dgrad, in1_dgrad, in2_dgrad, _ = buffers_copy
 
             weights_reordered = tp.reorder_weights_from_e3nn(
-                weights, not self.config.shared_weights
+                weights, not tp.config.shared_weights
             )
             weights_dgrad_reordered = tp.reorder_weights_from_e3nn(
-                weights_dgrad, not self.config.shared_weights
+                weights_dgrad, not tp.config.shared_weights
             )
 
-            in1_grad, in2_grad, weights_grad, out_dgrad = self.double_backward_cpu(in1, in2, out_grad, weights_reordered, weights_dgrad_reordered, in1_dgrad, in2_dgrad, graph)
+            in1_grad, in2_grad, weights_grad, out_dgrad = tp.double_backward_cpu(in1, in2, out_grad, weights_reordered, weights_dgrad_reordered, in1_dgrad, in2_dgrad, graph)
 
             tensors.append(
                 (   out_dgrad,
                     in1_grad,
                     in2_grad,
-                    self.reorder_weights_to_e3nn(weights_grad, has_batch_dim=not self.config.shared_weights) 
+                    tp.reorder_weights_to_e3nn(weights_grad, has_batch_dim=not self.config.shared_weights) 
                 )) 
 
         for name, to_check, ground_truth in [

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -165,10 +165,10 @@ def __call__(
         return self.forward(X, Y, W, rows, cols, sender_perm)
 
     def reorder_weights_from_e3nn(self, weights, has_batch_dim=True):
-        return reorder_jax(self.forward_schedule, weights, "forward", not self.config.shared_weights)
+        return reorder_jax(self.forward_schedule, weights, "forward", has_batch_dim) 
 
     def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
-        return reorder_jax(self.forward_schedule, weights, "backward", not self.config.shared_weights)
+        return reorder_jax(self.forward_schedule, weights, "backward", has_batch_dim)
 
     def forward_cpu(self, L1_in, L2_in, weights, L3_out, graph):
         rows = graph.rows.astype(np.int32)
@@ -240,4 +240,4 @@ def double_backward_cpu(self, in1, in2, out_grad, weights, weights_dgrad, in1_dg
             in1_jax, in2_jax, weights_jax, out_grad_jax
         )[1]((in1_dgrad_jax, in2_dgrad_jax, weights_dgrad_jax))
 
-        return in1_grad, in2_grad, weights_grad, out_dgrad
+        return np.asarray(in1_grad), np.asarray(in2_grad), np.asarray(weights_grad), np.asarray(out_dgrad) 

openequivariance/openequivariance/impl_torch/E3NNConv.py

@@ -37,7 +37,7 @@ def __init__(self, config, *, idx_dtype=np.int64, torch_op=True):
         if config.irrep_dtype == np.float64:
             torch.set_default_dtype(torch.float32)  # Reset to default
 
-    def forward(self, L1_in, L2_in, weights, rows, cols):
+    def forward(self, L1_in, L2_in, weights, rows, cols, transpose_perm=None):
         messages = self.reference_tp(L1_in[cols], L2_in, weights)
         return scatter_add_wrapper(messages, rows, L1_in.size(0))
 

openequivariance/pyproject.toml

@@ -48,8 +48,7 @@ bench = [
   "cuequivariance-ops-torch-cu12",
 ]
 
-jax = [
-  "jax[cuda12]",
+jax = [ 
   "nanobind",
   "scikit-build-core"
 ]