Linted.

Author: vbharadwaj-bk

openequivariance/openequivariance/__init__.py

@@ -31,19 +31,23 @@ def _check_package_editable():
 
 _editable_install_output_path = Path(__file__).parent.parent.parent / "outputs"
 
+
 def extension_source_path():
     """
     :returns: Path to the source code of the C++ extension.
     """
     return str(Path(__file__).parent / "extension")
 
+
 if "OEQ_NOTORCH" not in os.environ or os.environ["OEQ_NOTORCH"] != "1":
     import torch
 
-    from openequivariance._torch.TensorProduct import TensorProduct 
+    from openequivariance._torch.TensorProduct import TensorProduct
     from openequivariance._torch.TensorProductConv import TensorProductConv
 
-    from openequivariance._torch.extlib import torch_ext_so_path as torch_ext_so_path_internal
+    from openequivariance._torch.extlib import (
+        torch_ext_so_path as torch_ext_so_path_internal,
+    )
     from openequivariance.core.utils import torch_to_oeq_dtype
 
     torch.serialization.add_safe_globals(
@@ -60,6 +64,7 @@ def extension_source_path():
         ]
     )
 
+
 def torch_ext_so_path():
     """
     :returns: Path to a ``.so`` file that must be linked to use OpenEquivariance
@@ -70,6 +75,7 @@ def torch_ext_so_path():
     except NameError:
         return None
 
+
 jax = None
 try:
     import openequivariance_extjax
@@ -85,5 +91,5 @@ def torch_ext_so_path():
     "torch_to_oeq_dtype",
     "_check_package_editable",
     "torch_ext_so_path",
-    "jax"
+    "jax",
 ]

openequivariance/openequivariance/_torch/E3NNConv.py

@@ -7,6 +7,7 @@
 from openequivariance._torch.E3NNTensorProduct import E3NNTensorProduct
 from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
 
+
 class E3NNConv(ConvolutionBase, NumpyDoubleBackwardMixinConv):
     def __init__(self, config, *, idx_dtype=np.int64, torch_op=True):
         assert torch_op

openequivariance/openequivariance/_torch/NPDoubleBackwardMixin.py

@@ -1,75 +1,97 @@
 import torch
 
+
 class NumpyDoubleBackwardMixin:
-    '''
-    Adds a Numpy double backward method to any TensorProduct 
+    """
+    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):
+    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')
+        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
+            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]
+            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()
+        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):
+    """
+
+    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')
+        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)
+        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
+            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]
+            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()
-
-
+        return (
+            a.detach().cpu().numpy(),
+            b.detach().cpu().numpy(),
+            c.detach().cpu().numpy(),
+            d.detach().cpu().numpy(),
+        )

openequivariance/openequivariance/_torch/TensorProduct.py

@@ -92,10 +92,14 @@ def __setstate__(self, state):
         self._init_class()
 
     def reorder_weights_from_e3nn(self, weights, has_batch_dim=True):
-        return reorder_torch(self.forward_schedule, weights, "forward", not self.config.shared_weights)
+        return reorder_torch(
+            self.forward_schedule, weights, "forward", not self.config.shared_weights
+        )
 
     def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
-        return reorder_torch(self.forward_schedule, weights, "backward", not self.config.shared_weights)
+        return reorder_torch(
+            self.forward_schedule, weights, "backward", not self.config.shared_weights
+        )
 
     def forward(
         self, x: torch.Tensor, y: torch.Tensor, W: torch.Tensor
@@ -347,7 +351,7 @@ def name():
         return "LoopUnrollTP"
 
 
-if extlib.TORCH_COMPILE: 
+if extlib.TORCH_COMPILE:
     TensorProduct.register_torch_fakes()
     TensorProduct.register_autograd()
     TensorProduct.register_autocast()

openequivariance/openequivariance/_torch/TensorProductConv.py

@@ -27,6 +27,7 @@
 
 logger = getLogger()
 
+
 class TensorProductConv(torch.nn.Module, LoopUnrollConv, NumpyDoubleBackwardMixinConv):
     r"""
     Given a **symmetric, directed** graph :math:`G = (V, E)`, inputs :math:`x_1...x_{|V|}`,
@@ -420,10 +421,14 @@ def double_backward(ctx, grad_output):
         )
 
     def reorder_weights_from_e3nn(self, weights, has_batch_dim=True):
-        return reorder_torch(self.forward_schedule, weights, "forward", not self.config.shared_weights)
+        return reorder_torch(
+            self.forward_schedule, weights, "forward", not self.config.shared_weights
+        )
 
     def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
-        return reorder_torch(self.forward_schedule, weights, "backward", not self.config.shared_weights)
+        return reorder_torch(
+            self.forward_schedule, weights, "backward", not self.config.shared_weights
+        )
 
     @staticmethod
     def name():

openequivariance/openequivariance/_torch/utils.py

@@ -1,50 +1,53 @@
 import torch
 
+
 def reorder_helper(schedule, weights_in, direction, has_batch_dim):
-        assert direction in ["forward", "backward"]
-
-        specs = schedule.weight_reordering_info(weights_in, has_batch_dim) 
-        weights_out = torch.zeros_like(weights_in)
-
-        for spec in specs:
-            parent_range = spec["parent_range"]
-            parent_shape = spec["parent_shape"]
-            weights_subrange = spec["weights_subrange"]
-            child_range = spec["child_range"]
-            transpose_perm = spec["transpose_perm"]
-            
-            if direction == "forward":
-                reshape_size = spec["reshape_size"]
-                
-                sliced_weights = weights_in[parent_range].reshape(parent_shape)[
-                    weights_subrange
-                ]
-                
-                weights_out[child_range] = sliced_weights.permute(
-                    transpose_perm
-                ).reshape(reshape_size)
-
-            elif direction == "backward":
-                transpose_child_shape = spec["transpose_child_shape"]
-                child_shape = spec["child_shape"]
-
-                sliced_weights = (
-                    weights_in[child_range]
-                    .reshape(transpose_child_shape)
-                    .permute(transpose_perm)
-                )
-
-                weights_out[parent_range].reshape(parent_shape)[
-                    weights_subrange
-                ] = sliced_weights.flatten().reshape(child_shape)
-
-        return weights_out
+    assert direction in ["forward", "backward"]
+
+    specs = schedule.weight_reordering_info(weights_in, has_batch_dim)
+    weights_out = torch.zeros_like(weights_in)
+
+    for spec in specs:
+        parent_range = spec["parent_range"]
+        parent_shape = spec["parent_shape"]
+        weights_subrange = spec["weights_subrange"]
+        child_range = spec["child_range"]
+        transpose_perm = spec["transpose_perm"]
+
+        if direction == "forward":
+            reshape_size = spec["reshape_size"]
+
+            sliced_weights = weights_in[parent_range].reshape(parent_shape)[
+                weights_subrange
+            ]
+
+            weights_out[child_range] = sliced_weights.permute(transpose_perm).reshape(
+                reshape_size
+            )
+
+        elif direction == "backward":
+            transpose_child_shape = spec["transpose_child_shape"]
+            child_shape = spec["child_shape"]
+
+            sliced_weights = (
+                weights_in[child_range]
+                .reshape(transpose_child_shape)
+                .permute(transpose_perm)
+            )
+
+            weights_out[parent_range].reshape(parent_shape)[weights_subrange] = (
+                sliced_weights.flatten().reshape(child_shape)
+            )
+
+    return weights_out
+
 
 def reorder_numpy_helper(schedule, weights_in, direction, has_batch_dim):
     weights_in = torch.from_numpy(weights_in.copy())
     result = reorder_helper(schedule, weights_in, direction, has_batch_dim)
     return result.detach().cpu().numpy().copy()
 
+
 def reorder_torch(schedule, weights_in, direction, has_batch_dim):
     if isinstance(weights_in, torch.Tensor):
         return reorder_helper(schedule, weights_in, direction, has_batch_dim)

openequivariance/openequivariance/benchmark/correctness_utils.py

@@ -6,7 +6,8 @@
 from openequivariance.benchmark.random_buffer_utils import (
     get_random_buffers_forward,
     get_random_buffers_backward,
-    get_random_buffers_double_backward)
+    get_random_buffers_double_backward,
+)
 
 from openequivariance.benchmark.logging_utils import getLogger, bcolors
 import numpy as np
@@ -195,31 +196,51 @@ def correctness_double_backward(
     global torch
     import torch
 
-    in1, in2, out_grad, weights, weights_dgrad, in1_dgrad, in2_dgrad, _ = \
-        get_random_buffers_double_backward(problem, batch_size=batch_size, prng_seed=prng_seed) 
+    in1, in2, out_grad, weights, weights_dgrad, in1_dgrad, in2_dgrad, _ = (
+        get_random_buffers_double_backward(
+            problem, batch_size=batch_size, prng_seed=prng_seed
+        )
+    )
 
     if reference_implementation is None:
         from openequivariance._torch.E3NNTensorProduct import E3NNTensorProduct
+
         reference_implementation = E3NNTensorProduct
 
     result = {"thresh": correctness_threshold, "batch_size": batch_size}
 
     tensors = []
     for _, impl in enumerate([test_implementation, reference_implementation]):
         tp = instantiate_implementation(impl, problem)
-        weights_reordered = tp.reorder_weights_from_e3nn(weights, has_batch_dim=not problem.shared_weights)
-        weights_dgrad_reordered = tp.reorder_weights_from_e3nn(weights_dgrad, has_batch_dim=not problem.shared_weights)
+        weights_reordered = tp.reorder_weights_from_e3nn(
+            weights, has_batch_dim=not problem.shared_weights
+        )
+        weights_dgrad_reordered = tp.reorder_weights_from_e3nn(
+            weights_dgrad, has_batch_dim=not problem.shared_weights
+        )
 
         if impl == CUETensorProduct and problem.shared_weights:
             weights_reordered = weights_reordered[np.newaxis, :]
 
-        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)
+        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,
+        )
         tensors.append(
-            (   out_dgrad,
+            (
+                out_dgrad,
                 in1_grad,
                 in2_grad,
-                tp.reorder_weights_to_e3nn(weights_grad, has_batch_dim=not problem.shared_weights) 
-            )) 
+                tp.reorder_weights_to_e3nn(
+                    weights_grad, has_batch_dim=not problem.shared_weights
+                ),
+            )
+        )
 
     for name, to_check, ground_truth in [
         ("output_double_grad", tensors[0][0], tensors[1][0]),