Added the mixins.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_torch/E3NNConv.py

@@ -5,9 +5,9 @@
     scatter_add_wrapper,
 )
 from openequivariance.impl_torch.E3NNTensorProduct import E3NNTensorProduct
+from openequivariance.impl_torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
 
-
-class E3NNConv(ConvolutionBase):
+class E3NNConv(ConvolutionBase, NumpyDoubleBackwardMixinConv):
     def __init__(self, config, *, idx_dtype=np.int64, torch_op=True):
         assert torch_op
         super().__init__(config, idx_dtype=idx_dtype, torch_op=torch_op)

openequivariance/openequivariance/impl_torch/NPDoubleBackwardMixin.py

@@ -34,3 +34,42 @@ def double_backward_cpu(self, in1, in2, out_grad, weights, weights_dgrad, in1_dg
 
         return a.detach().cpu().numpy(), b.detach().cpu().numpy(), c.detach().cpu().numpy(), d.detach().cpu().numpy()
 
+
+class NumpyDoubleBackwardMixinConv:
+    '''
+    Similar, but for fused graph convolution.
+    '''
+    def double_backward_cpu(self, in1, in2, out_grad, weights, weights_dgrad, in1_dgrad, in2_dgrad, graph):
+        assert self.torch_op
+
+        in1_torch = torch.tensor(in1).to('cuda').requires_grad_(True)
+        in2_torch = torch.tensor(in2).to('cuda').requires_grad_(True)
+        weights_torch = torch.tensor(weights).to('cuda').requires_grad_(True)
+        out_grad_torch = torch.tensor(out_grad).to('cuda').requires_grad_(True)
+        in1_dgrad_torch = torch.tensor(in1_dgrad).to('cuda')
+        in2_dgrad_torch = torch.tensor(in2_dgrad).to('cuda')
+        weights_dgrad_torch = torch.tensor(weights_dgrad).to('cuda')
+
+        torch_rows = torch.tensor(graph.rows, device="cuda")
+        torch_cols = torch.tensor(graph.cols, device="cuda")
+        torch_transpose_perm = torch.tensor(graph.transpose_perm, device="cuda")
+
+        out_torch = self.forward(in1_torch, in2_torch, weights_torch, torch_rows, torch_cols, torch_transpose_perm)
+
+        in1_grad, in2_grad, weights_grad = torch.autograd.grad(
+            outputs=out_torch,
+            inputs=[in1_torch, in2_torch, weights_torch],
+            grad_outputs=out_grad_torch,
+            create_graph=True,
+            retain_graph=True
+        )
+
+        a, b, c, d = torch.autograd.grad(
+            outputs=[in1_grad, in2_grad, weights_grad],
+            inputs=[in1_torch, in2_torch, weights_torch, out_grad_torch],
+            grad_outputs=[in1_dgrad_torch, in2_dgrad_torch, weights_dgrad_torch]
+        )
+
+        return a.detach().cpu().numpy(), b.detach().cpu().numpy(), c.detach().cpu().numpy(), d.detach().cpu().numpy()
+
+

openequivariance/openequivariance/impl_torch/TensorProduct.py

@@ -7,7 +7,6 @@
 from openequivariance.benchmark.logging_utils import getLogger
 from openequivariance.impl_torch.utils import reorder_torch
 from openequivariance.impl_torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
-from openequivariance.core.e3nn_lite import Irreps
 
 logger = getLogger()
 
@@ -348,7 +347,7 @@ def name():
         return "LoopUnrollTP"
 
 
-if extlib.TORCH_COMPILE and __name__ != "__main__":
+if extlib.TORCH_COMPILE: 
     TensorProduct.register_torch_fakes()
     TensorProduct.register_autograd()
     TensorProduct.register_autocast()

openequivariance/openequivariance/impl_torch/TensorProductConv.py

@@ -22,11 +22,12 @@
 from openequivariance.impl_torch.utils import reorder_torch
 
 from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.impl_torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
 
-logger = getLogger()
 
+logger = getLogger()
 
-class TensorProductConv(torch.nn.Module, LoopUnrollConv):
+class TensorProductConv(torch.nn.Module, LoopUnrollConv, NumpyDoubleBackwardMixinConv):
     r"""
     Given a **symmetric, directed** graph :math:`G = (V, E)`, inputs :math:`x_1...x_{|V|}`,
     :math:`y_1...y_{|E|}`, and weights :math:`W_1...W_{|E|}`, computes