Torch scatter add (#143)

* replace custom scatter with torch scatter_add

* increase limit for backwards error

* move duplicate logic to an external function

* change to scatter_add_wrapper

---------

Co-authored-by: Austin Glover <austin_glover@berkeley.com>

Author: asglover

openequivariance/implementations/convolution/CUEConv.py

@@ -3,7 +3,10 @@
 from typing import Iterator
 
 from openequivariance.implementations.CUETensorProduct import CUETensorProduct
-from openequivariance.implementations.convolution.ConvolutionBase import ConvolutionBase
+from openequivariance.implementations.convolution.ConvolutionBase import (
+    ConvolutionBase,
+    scatter_add_wrapper,
+)
 
 
 class CUEConv(ConvolutionBase):
@@ -16,15 +19,9 @@ def __init__(self, config, *, idx_dtype=np.int64, torch_op=True):
         self.reference_tp = CUETensorProduct(config, torch_op)
         self.cue_tp = self.reference_tp.cue_tp
 
-        from openequivariance.implementations.convolution.scatter import scatter_sum
-
-        self.scatter_sum = scatter_sum
-
     def forward(self, L1_in, L2_in, weights, rows, cols):
-        tp_outputs = self.cue_tp(L1_in[cols], L2_in, weights)
-        return self.scatter_sum(
-            src=tp_outputs, index=rows, dim=0, dim_size=L1_in.shape[0]
-        )
+        messages = self.reference_tp(L1_in[cols], L2_in, weights)
+        return scatter_add_wrapper(messages, rows, L1_in.size(0))
 
     @staticmethod
     def name():

openequivariance/implementations/convolution/ConvolutionBase.py

@@ -713,7 +713,11 @@ def test_correctness_double_backward(
             in2_torch = torch.tensor(in2, device="cuda", requires_grad=True)
 
             weights_reordered = np.zeros_like(weights)
-            if i == 0 and self.reorder_weights_e3nn_to_oeq is not None:
+            if (
+                i == 0
+                and hasattr(self, "reorder_weights_e3nn_to_oeq")
+                and self.reorder_weights_e3nn_to_oeq is not None
+            ):
                 self.reorder_weights_e3nn_to_oeq(
                     weights, weights_reordered, not self.config.shared_weights
                 )
@@ -750,7 +754,11 @@ def test_correctness_double_backward(
             )
 
             weights_grad = weights_torch.grad.detach().cpu().numpy()
-            if i == 0 and self.reorder_weights_oeq_to_e3nn is not None:
+            if (
+                i == 0
+                and hasattr(self, "reorder_weights_e3nn_to_oeq")
+                and self.reorder_weights_oeq_to_e3nn is not None
+            ):
                 weights_grad_copy = weights_grad.copy()
                 self.reorder_weights_oeq_to_e3nn(
                     weights_grad_copy, weights_grad, not self.config.shared_weights
@@ -774,3 +782,15 @@ def test_correctness_double_backward(
             result[name] = check_similiarity(name, to_check, ground_truth, thresh)
 
         return result
+
+
+def scatter_add_wrapper(messages, rows, target_dim):
+    L3_dim = messages.size(1)
+    idx = rows.unsqueeze(1).expand(-1, L3_dim)
+    out = messages.new_zeros((target_dim, L3_dim))
+    return torch.scatter_add(
+        input=out,
+        dim=0,
+        index=idx,
+        src=messages,
+    )

openequivariance/implementations/convolution/E3NNConv.py

@@ -1,6 +1,9 @@
 import numpy as np
 
-from openequivariance.implementations.convolution.ConvolutionBase import ConvolutionBase
+from openequivariance.implementations.convolution.ConvolutionBase import (
+    ConvolutionBase,
+    scatter_add_wrapper,
+)
 from openequivariance.implementations.E3NNTensorProduct import E3NNTensorProduct
 
 
@@ -34,15 +37,9 @@ 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
 
-        from openequivariance.implementations.convolution.scatter import scatter_sum
-
-        self.scatter_sum = scatter_sum
-
     def forward(self, L1_in, L2_in, weights, rows, cols):
-        tp_outputs = self.reference_tp(L1_in[cols], L2_in, weights)
-        return self.scatter_sum(
-            src=tp_outputs, index=rows, dim=0, dim_size=L1_in.shape[0]
-        )
+        messages = self.reference_tp(L1_in[cols], L2_in, weights)
+        return scatter_add_wrapper(messages, rows, L1_in.size(0))
 
     @staticmethod
     def name():

openequivariance/implementations/convolution/TensorProductConv.py

@@ -4,7 +4,10 @@
 import torch
 
 from openequivariance import extlib
-from openequivariance.implementations.convolution.ConvolutionBase import ConvolutionBase
+from openequivariance.implementations.convolution.ConvolutionBase import (
+    ConvolutionBase,
+    scatter_add_wrapper,
+)
 from openequivariance.implementations.convolution.LoopUnrollConv import LoopUnrollConv
 from openequivariance.implementations.TensorProduct import TensorProduct
 from openequivariance import TPProblem
@@ -414,15 +417,12 @@ def __init__(self, config, *, torch_op=True):
         super().__init__(config, torch_op=torch_op, deterministic=False)
 
         self.reference_tp = TensorProduct(config, torch_op=torch_op)
-        from openequivariance.implementations.convolution.scatter import scatter_sum
-
-        self.scatter_sum = scatter_sum
+        self.reorder_weights_e3nn_to_oeq = self.reference_tp.reorder_weights_e3nn_to_oeq
+        self.reorder_weights_oeq_to_e3nn = self.reference_tp.reorder_weights_oeq_to_e3nn
 
     def forward(self, L1_in, L2_in, weights, rows, cols):
-        tp_outputs = self.reference_tp(L1_in[cols], L2_in, weights)
-        return self.scatter_sum(
-            src=tp_outputs, index=rows, dim=0, dim_size=L1_in.shape[0]
-        )
+        messages = self.reference_tp(L1_in[cols], L2_in, weights)
+        return scatter_add_wrapper(messages, rows, L1_in.size(0))
 
     def forward_cpu(self, L1_in, L2_in, weights, L3_out, graph):
         tp_outputs = np.zeros((graph.nnz, self.L3.dim), dtype=L3_out.dtype)

openequivariance/implementations/convolution/scatter.py

@@ -224,7 +224,7 @@ class TestAtomicSharedWeights(ConvCorrectness):
     def thresh(self, direction):
         return {
             "fwd": 1e-5,
-            "bwd": 5e-2,  # Expect higher errors for shared weights
+            "bwd": 7.5e-2,  # Expect higher errors for shared weights
             "double_bwd": 5e-2,
         }[direction]