Ready to modify the double backward correctness function.

Author: vbharadwaj-bk

openequivariance/openequivariance/benchmark/random_buffer_utils.py

@@ -104,10 +104,13 @@ def get_random_buffers_double_backward(
     )
     weights = np.array(rng.uniform(size=weights_size), dtype=tpp.irrep_dtype)
 
-    weights_grad = np.zeros_like(weights)
-    in1_grad = np.zeros_like(in1)
-    in2_grad = np.zeros_like(in2)
-    out_double_grad = np.zeros_like(out_grad)
+    weights_grad = np.array(rng.uniform(size=weights_size), dtype=tpp.irrep_dtype)
+    in1_grad = np.array(
+        rng.uniform(size=(batch_size, tpp.irreps_in1.dim)), dtype=tpp.irrep_dtype)
+    in2_grad = np.array(
+        rng.uniform(size=(batch_size, tpp.irreps_in2.dim)), dtype=tpp.irrep_dtype)
+    out_double_grad = np.array(
+        rng.uniform(size=(batch_size, tpp.irreps_out.dim)), dtype=tpp.irrep_dtype)
 
     return (
         in1,
@@ -176,3 +179,5 @@ def get_random_buffers_backward_conv(
     in2_grad = np.zeros_like(in2)
 
     return in1, in2, out_grad, weights, weights_grad, in1_grad, in2_grad
+
+

openequivariance/openequivariance/impl_torch/NPDoubleBackwardMixin.py

@@ -0,0 +1,36 @@
+import torch
+
+class NumpyDoubleBackwardMixin:
+    '''
+    Adds a Numpy double backward method to any TensorProduct 
+    with the forward pass defined in PyTorch and the relevant
+    derivatives registered. 
+    '''
+    def double_backward_cpu(self, in1, in2, out_grad, weights, weights_dgrad, in1_dgrad, in2_dgrad):
+        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')
+        out_torch = self.forward(in1_torch, in2_torch, weights_torch)
+
+        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

@@ -6,12 +6,13 @@
 from openequivariance.core.utils import torch_to_oeq_dtype
 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()
 
 
-class TensorProduct(torch.nn.Module, LoopUnrollTP):
+class TensorProduct(torch.nn.Module, LoopUnrollTP, NumpyDoubleBackwardMixin):
     r"""
     Drop-in replacement for ``o3.TensorProduct`` from e3nn. Supports forward,
     backward, and double-backward passes using JIT-compiled kernels. Initialization
@@ -347,7 +348,7 @@ def name():
         return "LoopUnrollTP"
 
 
-if extlib.TORCH_COMPILE:
+if extlib.TORCH_COMPILE and __name__ != "__main__":
     TensorProduct.register_torch_fakes()
     TensorProduct.register_autograd()
     TensorProduct.register_autocast()