Reformatted.

Author: vbharadwaj-bk

docs/conf.py

@@ -9,31 +9,30 @@
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
-project = 'OpenEquivariance'
-copyright = '2025, The Regents of the University of California, through Lawrence Berkeley National Laboratory.' 
-author = 'Vivek Bharadwaj, Austin Glover, Aydin Buluc, James Demmel'
+project = "OpenEquivariance"
+copyright = "2025, The Regents of the University of California, through Lawrence Berkeley National Laboratory."
+author = "Vivek Bharadwaj, Austin Glover, Aydin Buluc, James Demmel"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
 extensions = []
 
-templates_path = ['_templates']
-exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
-
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
-html_theme = 'furo'
-html_static_path = ['_static']
+html_theme = "furo"
+html_static_path = ["_static"]
 
 extensions = [
-    'sphinx.ext.autodoc',
+    "sphinx.ext.autodoc",
 ]
 
-sys.path.insert(0, str(Path('..').resolve()))
+sys.path.insert(0, str(Path("..").resolve()))
 
-autodoc_mock_imports = ['torch', 'openequivariance.extlib', 'jinja2'] 
+autodoc_mock_imports = ["torch", "openequivariance.extlib", "jinja2"]
 autodoc_typehints = "description"

openequivariance/__init__.py

@@ -11,12 +11,13 @@
 )
 from openequivariance.implementations.utils import torch_to_oeq_dtype
 
-__version__ = None 
+__version__ = None
 try:
     __version__ = version("openequivariance")
 except Exception as e:
     print(f"Warning: Could not determine oeq version: {e}", file=sys.stderr)
 
+
 def _check_package_editable():
     import json
     from importlib.metadata import Distribution
@@ -29,10 +30,10 @@ def _check_package_editable():
 
 
 def torch_ext_so_path():
-    '''
+    """
     :returns: Path to a ``.so`` file that must be linked to use OpenEquivariance
-              from the PyTorch C++ Interface. 
-    '''
+              from the PyTorch C++ Interface.
+    """
     return openequivariance.extlib.torch_module.__file__
 
 

openequivariance/implementations/TensorProduct.py

@@ -9,11 +9,12 @@ class TensorProduct(torch.nn.Module, LoopUnrollTP):
     backward, and double-backward passes using JIT-compiled kernels. Initialization
     fails if:
 
-    * There are no visible GPUs. 
-    * The provided tensor product specification is unsupported. 
-    
+    * There are no visible GPUs.
+    * The provided tensor product specification is unsupported.
+
     :param problem: Specification of the tensor product.
     """
+
     def __init__(self, problem: TPProblem, torch_op=True):
         torch.nn.Module.__init__(self)
         LoopUnrollTP.__init__(self, problem, torch_op)
@@ -26,19 +27,19 @@ def name():
     def forward(
         self, x: torch.Tensor, y: torch.Tensor, W: torch.Tensor
     ) -> torch.Tensor:
-        '''
+        """
         Computes :math:`W (x \otimes_{\\textrm{CG}} y)`, identical to
-        ``o3.TensorProduct.forward``. 
+        ``o3.TensorProduct.forward``.
 
         :param x: Tensor of shape ``[batch_size, problem.irreps_in1.dim()]``, datatype
                   ``problem.irrep_dtype``.
         :param y: Tensor of shape ``[batch_size, problem.irreps_in2.dim()]``, datatype
                   ``problem.irrep_dtype``.
-        :param W: Tensor of datatype ``problem.weight_dtype`` and shape 
+        :param W: Tensor of datatype ``problem.weight_dtype`` and shape
 
             * ``[batch_size, problem.weight_numel]`` if ``problem.shared_weights=False``
             * ``[problem.weight_numel]`` if ``problem.shared_weights=True``
 
         :return: Tensor of shape ``[batch_size, problem.irreps_out.dim()]``, datatype ``problem.irrep_dtype``.
-        '''
+        """
         return torch.ops.libtorch_tp_jit.jit_tp_forward(self.internal, x, y, W)

openequivariance/implementations/convolution/TensorProductConv.py

@@ -10,34 +10,35 @@
 from openequivariance.implementations.TensorProduct import TensorProduct
 from openequivariance import TPProblem
 
+
 class TensorProductConv(torch.nn.Module, LoopUnrollConv):
     """
-    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 
+    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
 
     .. math::
         z_i = \sum_{(i, j, e) \in \mathcal{N}(i)} W_e (x_j \otimes_{\\textrm{CG}} y_e)
-    where :math:`(i, j, e) \in \mathcal{N}(i)` indicates that node :math:`i` is connected to node :math:`j` 
-    via the edge indexed :math:`e`. 
+    where :math:`(i, j, e) \in \mathcal{N}(i)` indicates that node :math:`i` is connected to node :math:`j`
+    via the edge indexed :math:`e`.
 
     This class offers multiple options to perform the summation: an atomic algorithm and a deterministic algorithm
     that relies on a sorted adjacency matrix input. If you use the determinstic algorithm, you must also supply
-    a permutation to transpose the adjacency matrix. 
+    a permutation to transpose the adjacency matrix.
 
     :param problem: Specification of the tensor product.
     :param deterministic: if ``False``, uses atomics for the convolution. If ``True``, uses a deterministic
            fixup-based algorithm. `Default`: ``False``.
     :param kahan: if ``True``, uses Kahan summation to improve accuracy during aggregation. To use this option,
-              the input tensors must be in float32 precision AND you must set ``deterministic=True``. *Default*: ``False``. 
-    
+              the input tensors must be in float32 precision AND you must set ``deterministic=True``. *Default*: ``False``.
+
     """
 
     def __init__(
         self,
         problem: TPProblem,
-        deterministic: bool=False,
-        kahan: bool =False,
-        torch_op=True
+        deterministic: bool = False,
+        kahan: bool = False,
+        torch_op=True,
     ):
         torch.nn.Module.__init__(self)
         LoopUnrollConv.__init__(
@@ -64,26 +65,26 @@ def forward(
         cols: torch.Tensor,
         sender_perm: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        '''
-        Computes the fused CG tensor product + convolution. 
+        """
+        Computes the fused CG tensor product + convolution.
 
         :param X: Tensor of shape ``[|V|, problem.irreps_in1.dim()]``, datatype ``problem.irrep_dtype``.
         :param Y: Tensor of shape ``[|E|, problem.irreps_in1.dim()]``, datatype ``problem.irrep_dtype``.
-        :param W: Tensor of datatype ``problem.weight_dtype`` and shape 
+        :param W: Tensor of datatype ``problem.weight_dtype`` and shape
 
             * ``[|E|, problem.weight_numel]`` if ``problem.shared_weights=False``
             * ``[problem.weight_numel]`` if ``problem.shared_weights=True``
-        
-        :param rows: Tensor of shape ``[|E|]`` with row indices for each nonzero in the adjacency matrix, 
+
+        :param rows: Tensor of shape ``[|E|]`` with row indices for each nonzero in the adjacency matrix,
                 datatype ``torch.int64``. Must be row-major sorted along with ``cols`` when ``deterministic=True``.
         :param cols: Tensor of shape ``[|E|]`` with column indices for each nonzero in the adjacency matrix,
-                datatype ``torch.int64``. 
-        :param sender_perm: Tensor of shape ``[|E|]`` and ``torch.int64`` datatype containing a 
+                datatype ``torch.int64``.
+        :param sender_perm: Tensor of shape ``[|E|]`` and ``torch.int64`` datatype containing a
                 permutation that transposes the adjacency matrix nonzeros from row-major to column-major order.
                 Must be provided when ``deterministic=True``.
-        
+
         :return: Tensor of shape ``[|V|, problem.irreps_out.dim()]``, datatype ``problem.irrep_dtype``.
-        '''
+        """
         if sender_perm is None:
             return torch.ops.libtorch_tp_jit.jit_conv_forward(
                 self.internal,

openequivariance/implementations/e3nn_lite.py

@@ -359,23 +359,23 @@ class Instruction(NamedTuple):
 
 
 class TPProblem:
-    '''
-    Specification for a CG tensor product. All parameters from 
+    """
+    Specification for a CG tensor product. All parameters from
     e3nn's ``o3.TensorProduct`` are available, along with additional
     parameters for the types of weights and irreps.
 
     :param irreps_in1: Irreps for the first CG argument
     :param irreps_in2: Irreps for the second CG argument
-    :param irreps_out: Irreps for the output 
-    :param instructions: A list of 5-tuples, each of 
+    :param irreps_out: Irreps for the output
+    :param instructions: A list of 5-tuples, each of
            the form ``(i_in1, i_in2, i_out, has_weight, path_weight)``.
            ``i_in1``, ``i_in2``, and ``i_out`` each index
            an Irrep from ``irreps_in1``, ``irreps_in2``, and
            ``irreps_in3``, respectively. ``has_weight`` (True / False)
            controls whether trainable weights are included for the
-           instruction, and ``path_weight`` controls output normalization. 
+           instruction, and ``path_weight`` controls output normalization.
     :param irrep_dtype: Datatype of irrep inputs; one of ``np.float32`` or ``np.float64``.
-           *Default*: ``np.float32``. 
+           *Default*: ``np.float32``.
     :param weight_dtype: Datatype of weights; one of ``np.float32`` or ``np.float64``.
            *Default*: ``np.float32``.
     :param label: A name for this problem specification (useful for testing / benchmarking).
@@ -384,7 +384,8 @@ class TPProblem:
     :param internal_weights: Must be False; OpenEquivariance does not support internal weights. *Default*: False.
     :param irrep_normalization: One of ``["component", "norm", "none"]``. *Default*: "component".
     :param path_normalization: One of ``["element", "path", "none"]``. *Default*: "element".
-    '''
+    """
+
     instructions: List[Any]
     shared_weights: bool
     internal_weights: bool

openequivariance/implementations/utils.py

@@ -79,7 +79,7 @@ def filter_and_analyze_problem(problem):
         f"irrep_dtype and weight_dtype must be the same, got {problem.irrep_dtype} and {problem.weight_dtype}"
     )
 
-    assert not problem.internal_weights, ( 
+    assert not problem.internal_weights, (
         f"Openequivariance does not support internal weights, got {problem.internal_weights}"
     )