Ran ruff.

Author: vbharadwaj-bk

docs/conf.py

@@ -34,5 +34,10 @@
 
 sys.path.insert(0, str(Path("..").resolve()))
 
-autodoc_mock_imports = ["torch", "openequivariance.impl_torch.extlib", "jinja2", "numpy"]
+autodoc_mock_imports = [
+    "torch",
+    "openequivariance.impl_torch.extlib",
+    "jinja2",
+    "numpy",
+]
 autodoc_typehints = "description"

openequivariance/openequivariance/__init__.py

@@ -17,6 +17,7 @@
 from openequivariance.impl_torch.TensorProductConv import (
     TensorProductConv,
 )
+from openequivariance.impl_torch.extlib import torch_ext_so_path
 from openequivariance.core.utils import torch_to_oeq_dtype
 
 __version__ = None
@@ -37,20 +38,13 @@ def _check_package_editable():
 _editable_install_output_path = Path(__file__).parent.parent.parent / "outputs"
 
 
-def torch_ext_so_path():
-    """
-    :returns: Path to a ``.so`` file that must be linked to use OpenEquivariance
-              from the PyTorch C++ Interface.
-    """
-    return openequivariance.impl_torch.extlib.torch_module.__file__
-
-
 def extension_source_path():
     """
     :returns: Path to the source code of the C++ extension.
     """
     return str(Path(__file__).parent / "extension")
 
+
 torch.serialization.add_safe_globals(
     [
         TensorProduct,

openequivariance/openequivariance/core/ConvolutionBase.py

@@ -562,6 +562,7 @@ def test_correctness_double_backward(
 
         if reference_implementation is None:
             from openequivariance.impl_torch.E3NNConv import E3NNConv
+
             reference_implementation = E3NNConv
 
         reference_problem = self.config

openequivariance/openequivariance/core/LoopUnrollConv.py

@@ -10,11 +10,13 @@
 from openequivariance.templates.jinja_utils import get_jinja_environment
 from openequivariance.core.utils import filter_and_analyze_problem
 
+
 class LoopUnrollConv(ConvolutionBase):
     def __init__(
         self,
         config,
-        dp, postprocess_kernel,
+        dp,
+        postprocess_kernel,
         *,
         idx_dtype: type[np.generic] = np.int64,
         torch_op: bool = False,

openequivariance/openequivariance/core/LoopUnrollTP.py

@@ -10,6 +10,7 @@
     count_cg_non_zero,
 )
 
+
 class LoopUnrollTP(TensorProductBase):
     def __init__(self, config, dp, postprocess_kernel, torch_op):
         super().__init__(config, torch_op=torch_op)
@@ -99,14 +100,12 @@ def generate_double_backward_schedule(warps_per_block):
             "opt_level": 3,
             "irrep_dtype": dtype_to_enum[self.config.irrep_dtype],
             "weight_dtype": dtype_to_enum[self.config.weight_dtype],
-
-            # Not relevant, included for compatibility with convolution 
+            # Not relevant, included for compatibility with convolution
             "workspace_size": 0,
             "deterministic": 1,
-            "idx_dtype": 0
+            "idx_dtype": 0,
         }
 
-
     def calculate_flops_forward(self, batch_size: int) -> dict:
         if self.is_uvw:
             return super().calculate_flops_forward(batch_size)

openequivariance/openequivariance/core/utils.py

@@ -10,6 +10,7 @@
 import hashlib
 from openequivariance.impl_torch.extlib import GPUTimer
 
+
 def sparse_outer_product_work(cg: np.ndarray) -> int:
     return np.sum(np.max(cg != 0, axis=2))
 

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -1,54 +1,62 @@
-import numpy as np
-
 import jax
-
 from functools import partial
 from openequivariance.impl_jax import extlib
-import hashlib
-from openequivariance.core.e3nn_lite import TPProblem, Irreps
+from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.core.LoopUnrollTP import LoopUnrollTP
 from openequivariance.core.utils import hash_attributes
-import jax.numpy as jnp
 
-@partial(jax.custom_vjp, nondiff_argnums=(3,4,5))
+
+@partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5))
 def forward(X, Y, W, L3_dim, irrep_dtype, attrs):
-    forward_call = jax.ffi.ffi_call("tp_forward", 
-        jax.ShapeDtypeStruct((X.shape[0], L3_dim), irrep_dtype))
+    forward_call = jax.ffi.ffi_call(
+        "tp_forward", jax.ShapeDtypeStruct((X.shape[0], L3_dim), irrep_dtype)
+    )
     return forward_call(X, Y, W, **attrs)
 
+
 def forward_with_inputs(X, Y, W, L3_dim, irrep_dtype, attrs):
     return forward(X, Y, W, L3_dim, irrep_dtype, attrs), (X, Y, W)
 
-@partial(jax.custom_vjp, nondiff_argnums=(4,5))
+
+@partial(jax.custom_vjp, nondiff_argnums=(4, 5))
 def backward(X, Y, W, dZ, irrep_dtype, attrs):
-    backward_call = jax.ffi.ffi_call("tp_backward",
+    backward_call = jax.ffi.ffi_call(
+        "tp_backward",
         (
             jax.ShapeDtypeStruct(X.shape, irrep_dtype),
             jax.ShapeDtypeStruct(Y.shape, irrep_dtype),
             jax.ShapeDtypeStruct(W.shape, irrep_dtype),
-        ))
+        ),
+    )
 
     return backward_call(X, Y, W, dZ, **attrs)
 
+
 def backward_with_inputs(X, Y, W, dZ, irrep_dtype, attrs):
     return backward(X, Y, W, dZ, irrep_dtype, attrs), (X, Y, W, dZ)
 
+
 def double_backward(irrep_dtype, attrs, inputs, derivatives):
-    double_backward_call = jax.ffi.ffi_call("tp_double_backward",
+    double_backward_call = jax.ffi.ffi_call(
+        "tp_double_backward",
         (
             jax.ShapeDtypeStruct(inputs[0].shape, irrep_dtype),
             jax.ShapeDtypeStruct(inputs[1].shape, irrep_dtype),
             jax.ShapeDtypeStruct(inputs[2].shape, irrep_dtype),
             jax.ShapeDtypeStruct(inputs[3].shape, irrep_dtype),
-        ))
+        ),
+    )
     return double_backward_call(*inputs, *derivatives, **attrs)
 
+
 def backward_autograd(L3_dim, irrep_dtype, attrs, inputs, dZ):
-    return backward(inputs[0], inputs[1], inputs[2], dZ, irrep_dtype, attrs) 
+    return backward(inputs[0], inputs[1], inputs[2], dZ, irrep_dtype, attrs)
+
 
 forward.defvjp(forward_with_inputs, backward_autograd)
 backward.defvjp(backward_with_inputs, double_backward)
 
+
 class TensorProduct(LoopUnrollTP):
     def __init__(self, config: TPProblem):
         dp = extlib.DeviceProp(0)
@@ -59,18 +67,17 @@ def __init__(self, config: TPProblem):
             "forward_config": vars(self.forward_schedule.launch_config),
             "backward_config": vars(self.backward_schedule.launch_config),
             "double_backward_config": vars(self.double_backward_schedule.launch_config),
-            "kernel_prop": self.kernelProp
+            "kernel_prop": self.kernelProp,
         }
         hash_attributes(self.attrs)
- 
+
         self.weight_numel = config.weight_numel
         self.L3_dim = self.config.irreps_out.dim
 
     def forward(self, X: jax.ndarray, Y: jax.ndarray, W: jax.ndarray) -> jax.ndarray:
         return forward(X, Y, W, self.L3_dim, self.config.irrep_dtype, self.attrs)
 
-    def __call__(self,
-            X: jax.numpy.ndarray, 
-            Y: jax.numpy.ndarray, 
-            W: jax.numpy.ndarray) -> jax.numpy.ndarray:
-        return self.forward(X, Y, W)
+    def __call__(
+        self, X: jax.numpy.ndarray, Y: jax.numpy.ndarray, W: jax.numpy.ndarray
+    ) -> jax.numpy.ndarray:
+        return self.forward(X, Y, W)