Fix pylint mypy issues

cirkit/backend/torch/compiler.py

@@ -448,7 +448,8 @@ def _fold_parameters(
 def _fold_parameter_nodes_group(
     group: list[TorchParameterNode], *, compiler: TorchCompiler
 ) -> TorchParameterNode:
-    """Fold a list of [TorchParameterNode][cirkit.backend.torch.parameters.nodes.TorchParameterNode].
+    """Fold a list of
+    [TorchParameterNode][cirkit.backend.torch.parameters.nodes.TorchParameterNode].
 
     Args:
         group (list[TorchParameterNode]): List of parameter to fold.

cirkit/backend/torch/graph/optimize.py

@@ -48,9 +48,11 @@ def entries(cls) -> Sequence[type[TorchModuleT]]:
     def config_patterns(cls) -> Sequence[Mapping[str, Any]]:
         """Returns a list of dictionaries that match a config name to a config value.
 
-        The “config” of a layer / parameter node is simply the dictionary returned by `Layer.config`.
+        The “config” of a layer / parameter node is simply the dictionary returned
+        by `Layer.config`.
 
-        The dictionary at position x in the list define the config for the x-th element of the `entries` list.
+        The dictionary at position x in the list define the config for the x-th element
+        of the `entries` list.
 
         For example, the sum layer with config:
 
@@ -80,11 +82,14 @@ def entries():
 
     @classmethod
     def sub_patterns(cls) -> Sequence[Mapping[str, type["GraphOptPatternDefn"]]]:
-        """Returns a list of dictionaries that map layer's parameter names to a `ParameterOptPattern`.
+        """Returns a list of dictionaries that map layer's parameter names to a
+        `ParameterOptPattern`.
 
-        The dictionary at position x in the list define the config for the x-th element of the `entries` list.
+        The dictionary at position x in the list define the config for the x-th element
+        of the `entries` list.
 
-        For example, you can match the weight parameter of a sum layer to be of a certain `ParameterType`:
+        For example, you can match the weight parameter of a sum layer to be of a certain
+        `ParameterType`:
 
         ```python
         class LayerPatternOne(LayerOptPatternDefn):
@@ -357,7 +362,8 @@ def match_optimization_patterns(
     Returns:
         tuple[list[GraphOptMatch[TorchModuleT]], dict[TorchModuleT, GraphOptMatch[TorchModuleT]]]:
             - List of all the matches after priority-based  filtering
-            - Mapping from the module that matches the root of the pattern to the corresponding match.
+            - Mapping from the module that matches the root of the pattern to the
+              corresponding match.
     """
     ordering = list(ordering) if isinstance(ordering, Iterator) else ordering
     outputs = list(outputs) if isinstance(outputs, Iterator) else outputs
@@ -399,21 +405,24 @@ def _prioritize_optimization_strategy(
     in_place: bool = True,
 ) -> dict[TorchModuleT, GraphOptMatch[TorchModuleT]]:
     """Sort matches according to an optimization strategy and select the highest priority.
-    An optimization pattern match is currently a composition of torch modules, i.e., part of the circuit or a part of the parameter computational graph.
+    An optimization pattern match is currently a composition of torch modules, i.e.,
+    part of the circuit or a part of the parameter computational graph.
     Each torch module can take part in one or more optimization pattern matches.
-    This function resolves such optimization conflicts by using a given optimization prioritization strategy.
+    This function resolves such optimization conflicts by using a given optimization
+    prioritization strategy.
 
     Args:
         ordering (Iterable[TorchModuleT]): List of compiled module in the graph.
-        module_matches (dict[TorchModuleT, list[GraphOptMatch[TorchModuleT]]]): mapping between modules
-            and list of matches.
+        module_matches (dict[TorchModuleT, list[GraphOptMatch[TorchModuleT]]]):
+            mapping between modules and list of matches.
         strategy (OptMatchStrategy, optional): The strategy to use when sorting the matches.
              Defaults to OptMatchStrategy.LARGEST_MATCH.
         in_place (bool, optional): if True, the function will directly modify `module_matches`.
              Defaults to True.
 
     Returns:
-        dict[TorchModuleT, GraphOptMatch[TorchModuleT]]: mapping between modules and exactly one match.
+        dict[TorchModuleT, GraphOptMatch[TorchModuleT]]: mapping between modules
+            and exactly one match.
     """
 
     if not in_place:
@@ -440,14 +449,16 @@ def _prioritize_optimization_strategy(
             )
             if prioritized_match_optional is not None:
                 # If an optimization pattern match has already been selected, then we select it,
-                # and mark all the other optimization pattern matches involving the module to be removed next
+                # and mark all the other optimization pattern matches involving the
+                # module to be removed next
                 prioritized_match = prioritized_match_optional
                 remaining_matches = [m for m in matches if m is not prioritized_match]
             else:
                 # Otherwise, we sort the matches based on the given strategy
                 # We select the optimization pattern match having the highest
                 # priority according to the given strategy (e.g., largest pattern)
-                # The remaining optimization pattern matches involving the module are then marked to be removed
+                # The remaining optimization pattern matches involving the module are
+                # then marked to be removed
                 sorted_matches = _sort_matches_priority(matches, strategy=strategy)
                 prioritized_match = sorted_matches[0]
                 remaining_matches = sorted_matches[1:]

cirkit/backend/torch/optimization/layers.py

@@ -286,7 +286,8 @@ def _apply_tensordot_rule(
     weight: TorchParameter,
     kronecker: TorchKroneckerParameter,
 ) -> tuple[TorchTensorDotLayer, TorchTensorDotLayer]:
-    r"""Returns the two dot layer representing a sum parameterized by the output of a Kronecker Product.
+    r"""Returns the two dot layer representing a sum parameterized by the output of a
+    Kronecker Product.
 
     This trick comes from (Zhang et al., 2025) Subsection 3.1.
     Given $W=A \otimes B$ the parameters of a sum or dot layer,

cirkit/backend/torch/optimization/parameters.py

@@ -55,6 +55,8 @@ def entries(cls) -> list[type[TorchParameterNode]]:
 
 
 class ReduceSumOuterProductPattern(ParameterOptPatternDefn):
+    """Detect a reduce-sum node applied to the result of an outer-product node."""
+
     @classmethod
     def is_output(cls) -> bool:
         return False

cirkit/backend/torch/parameters/nodes.py

@@ -221,7 +221,8 @@ def forward(self) -> Tensor:
 
 
 class TorchPointerParameter(TorchParameterInput):
-    """Reprensents fold based slices of an existing [TorchTensorParameter][cirkit.backend.torch.parameters.nodes.TorchTensorParameter].
+    """Reprensents fold based slices of an existing
+    [TorchTensorParameter][cirkit.backend.torch.parameters.nodes.TorchTensorParameter].
     These slices can be:
     - A single fold index.
     - A list of potentially non contiguous fold index.
@@ -736,7 +737,7 @@ class TorchSoftplusParameter(TorchEntrywiseParameterOp):
     """
 
     def forward(self, x: Tensor) -> Tensor:
-        return torch.nn.functional.softplus(x)
+        return torch.nn.functional.softplus(x)  # pylint: disable=not-callable
 
 
 class TorchConjugateParameter(TorchEntrywiseParameterOp):

cirkit/backend/torch/parameters/pic.py

@@ -178,8 +178,8 @@ def __repr__(self):
         )
 
 
-class PICInnerNet(nn.Module):
-    def __init__(
+class PICInnerNet(nn.Module):  # pylint: disable=too-many-instance-attributes
+    def __init__(  # pylint: disable=too-many-arguments,too-many-positional-arguments
         self,
         num_dim: int,
         num_funcs: int,
@@ -259,7 +259,7 @@ def _register_forward_hook(self, tensor_parameter: TorchTensorParameter) -> None
         """Register a forward hook on the tensor parameter to call this PICInnerNet."""
         pic_net = self  # Capture reference to self
 
-        def forward_hook(module, input, output):
+        def forward_hook(_module, _input, _output):
             # Replace the output with the PICInnerNet's output
             return pic_net()
 
@@ -333,7 +333,7 @@ def _is_mixing_weight_tensor(
 
 
 @torch.no_grad()
-def pc2qpc(
+def pc2qpc(  # pylint: disable=too-many-locals,too-many-branches,protected-access
     pc: TorchCircuit,
     integration_method: str,
     net_dim: int | None = 128,

cirkit/backend/torch/semiring.py

@@ -13,6 +13,8 @@
 
 
 class EinsumFunc(Protocol):
+    """Protocol for einsum-like tensor-contraction callables."""
+
     def __call__(self, *xs: Tensor) -> Tensor: ...
 
 

cirkit/backend/torch/utils.py

@@ -8,6 +8,12 @@
 
 # pylint: disable-next=abstract-method
 class SafeLog(autograd.Function):
+    """A numerically safe natural logarithm autograd function.
+
+    In the backward pass, it replaces the NaNs and infinities
+    arising in the gradient with zeros.
+    """
+
     @staticmethod
     def forward(x: Tensor) -> Tensor:  # pylint: disable=arguments-differ
         return torch.log(x)
@@ -30,6 +36,12 @@ def backward(ctx: Any, grad_output: Tensor) -> Tensor:  # pylint: disable=argume
 
 # pylint: disable-next=abstract-method
 class ComplexSafeLog(autograd.Function):
+    """A numerically safe natural logarithm autograd function for complex inputs.
+
+    In the backward pass, it replaces the NaNs and infinities
+    arising in the gradient with zeros.
+    """
+
     @staticmethod
     def forward(x: Tensor) -> Tensor:  # pylint: disable=arguments-differ
         return torch.log(x)

cirkit/symbolic/io.py

@@ -8,7 +8,7 @@
 from cirkit.symbolic.layers import HadamardLayer, InputLayer, KroneckerLayer, ProductLayer, SumLayer
 
 
-def plot_circuit(
+def plot_circuit(  # pylint: disable=too-many-arguments,too-many-positional-arguments
     circuit: Circuit,
     out_path: str | PathLike[str] | None = None,
     orientation: str = "vertical",

cirkit/templates/logic/graph.py

@@ -1,8 +1,7 @@
 import itertools
 from abc import ABC
-from collections.abc import Iterator, Sequence
+from collections.abc import Iterator, Mapping, Sequence
 from functools import cached_property
-from typing import cast
 
 from cirkit.symbolic.circuit import Circuit
 from cirkit.symbolic.initializers import ConstantTensorInitializer
@@ -73,10 +72,12 @@ class DisjunctionNode(LogicalCircuitNode):
 
 
 class LogicalCircuit(RootedDiAcyclicGraph[LogicalCircuitNode]):
+    """A logic circuit represented as a rooted directed acyclic graph."""
+
     def __init__(
         self,
         nodes: Sequence[LogicalCircuitNode],
-        in_nodes: dict[LogicalCircuitNode, Sequence[LogicalCircuitNode]],
+        in_nodes: Mapping[LogicalCircuitNode, Sequence[LogicalCircuitNode]],
         outputs: Sequence[LogicalCircuitNode],
     ) -> None:
         """A Logical circuit represented as a rooted acyclic graph.
@@ -92,7 +93,7 @@ def __init__(
             assert ValueError("A logic graphs can only have one output!")
         super().__init__(nodes, in_nodes, outputs)
 
-    def prune(self):
+    def prune(self) -> None:
         """Prune the current graph by applying unit propagation.
 
         Prune a graph in place by applying unit propagation to conjunction and disjunctions.
@@ -103,7 +104,7 @@ def prune(self):
         absorbing_element = lambda n: BottomNode if isinstance(n, ConjunctionNode) else TopNode
         null_element = lambda n: TopNode if isinstance(n, ConjunctionNode) else BottomNode
 
-        def absorb_node(node):
+        def absorb_node(node: LogicalCircuitNode) -> LogicalCircuitNode:
             if isinstance(node, (ConjunctionNode, DisjunctionNode)):
                 children = [absorb_node(c) for c in self.node_inputs(node)]
 
@@ -115,7 +116,7 @@ def absorb_node(node):
             return node
 
         # apply node absorbion and remove null elements from conjunctions and disjunctions
-        in_nodes = {}
+        in_nodes: dict[LogicalCircuitNode, list[LogicalCircuitNode]] = {}
         for n, children in self._in_nodes.items():
             absorbed = absorb_node(n)
 
@@ -137,18 +138,23 @@ def absorb_node(node):
         nodes = list(set(itertools.chain(*in_nodes.values())).union(in_nodes.keys()))
 
         # re initialize the graph
-        self.__init__(nodes, in_nodes, list(self.outputs))
+        super().__init__(nodes, in_nodes, list(self.outputs))
 
     @property
     def inputs(self) -> Iterator[LogicalCircuitNode]:
-        return (cast(LogicalCircuitNode, node) for node in super().inputs)
+        return super().inputs
 
     @property
-    def outputs(self) -> Iterator[LogicalCircuitNode]:
-        return (cast(LogicalCircuitNode, node) for node in super().outputs)
+    def outputs(self) -> Sequence[LogicalCircuitNode]:
+        return super().outputs
 
     @cached_property
     def num_variables(self) -> int:
+        """The number of distinct variables in the logic circuit.
+
+        Returns:
+            int: The number of distinct literals appearing in the input nodes.
+        """
         return len({i.literal for i in self.inputs if isinstance(i, LogicalInputNode)})
 
     def node_scope(self, node: LogicalCircuitNode) -> Scope:
@@ -174,14 +180,11 @@ def node_scope(self, node: LogicalCircuitNode) -> Scope:
 
         return scope
 
-    def smooth(self) -> "LogicalCircuit":
+    def smooth(self) -> None:
         """Convert the current graph to a smooth graph in place.
         see https://yoojungchoi.github.io/files/ProbCirc20.pdf and
         https://proceedings.neurips.cc/paper/2019/file/940392f5f32a7ade1cc201767cf83e31-Paper.pdf
         for more information.
-
-        Returns:
-            LogicalCircuit: A new logic graph that is smooth.
         """
         literal_map: dict[tuple[int, bool], LogicalCircuitNode] = {
             (node.literal, isinstance(node, LiteralNode)): node
@@ -192,7 +195,9 @@ def smooth(self) -> "LogicalCircuit":
         smoothing_map: dict[int, DisjunctionNode] = {}
         disjunctions = [n for n in self.nodes if isinstance(n, DisjunctionNode)]
 
-        in_nodes = self._in_nodes
+        in_nodes: dict[LogicalCircuitNode, list[LogicalCircuitNode]] = {
+            node: list(inputs) for node, inputs in self.nodes_inputs.items()
+        }
         for d in disjunctions:
             d_scope = self.node_scope(d)
 
@@ -229,7 +234,7 @@ def smooth(self) -> "LogicalCircuit":
                         in_nodes[d].insert(0, ad_hoc)
 
         nodes = list(set(itertools.chain(*in_nodes.values())).union(in_nodes.keys()))
-        self.__init__(nodes, in_nodes, self._outputs)
+        super().__init__(nodes, in_nodes, self._outputs)
 
     def build_circuit(
         self,
@@ -279,13 +284,15 @@ def build_circuit(
             # default factory is locally imported when needed to avoid circular imports
             literal_input_factory = default_literal_input_factory(negated=False)
             negated_literal_input_factory = default_literal_input_factory(negated=True)
+        assert literal_input_factory is not None and negated_literal_input_factory is not None
 
         if weight_factory is None:
             # default to unitary weights
-            def weight_factory(n: tuple[int]) -> Parameter:
-                # locally import numpy to avoid dependency on the whole file
+            def default_weight_factory(shape: tuple[int, ...]) -> Parameter:
                 initializer = ConstantTensorInitializer(1.0)
-                return Parameter.from_input(TensorParameter(*n, initializer=initializer))
+                return Parameter.from_input(TensorParameter(*shape, initializer=initializer))
+
+            weight_factory = default_weight_factory
 
         # map each input literal to a symbolic input layer
         for i in self.inputs:

cirkit/templates/logic/sdd.py

@@ -15,6 +15,8 @@
 
 
 class SDD(LogicalCircuit):
+    """A Sentential Decision Diagram (SDD) represented as a logical circuit."""
+
     @staticmethod
     def load(filename: str) -> "SDD":
         """Load the SDD from a file.
@@ -53,7 +55,7 @@ def load(filename: str) -> "SDD":
                         n_id, _, l = args
 
                         if l > 0:
-                            node = LiteralNode(abs(l) - 1)
+                            node: LogicalCircuitNode = LiteralNode(abs(l) - 1)
                             nodes_map[n_id] = node
                             literal_map[(abs(l), True)] = node
                         else:

cirkit/templates/region_graph/algorithms/chow_liu.py

@@ -1,4 +1,5 @@
 from warnings import warn
+
 import numpy as np
 import torch
 from scipy import sparse as sp

cirkit/templates/region_graph/algorithms/quad.py

@@ -17,7 +17,9 @@ def QuadTree(shape: tuple[int, int, int], *, num_patch_splits: int = 2) -> Regio
     r"""Constructs a Quad Tree region graph.
 
     See:
-        - *Unifying and understanding overparameterized circuit representations via low-rank tensor decompositions.* [🔗](https://openreview.net/forum?id=1btutFdIya)
+        - *Unifying and understanding overparameterized circuit representations via
+          low-rank tensor decompositions.*
+          [🔗](https://openreview.net/forum?id=1btutFdIya)
           Mari, Antonio, Gennaro Vessio, and Antonio Vergari.
           In The 6th Workshop on Tractable Probabilistic Modeling. 2023.
 
@@ -41,7 +43,9 @@ def QuadGraph(shape: tuple[int, int, int]) -> RegionGraph:
     r"""Constructs a Quad Graph region graph.
 
     See:
-        - *Unifying and understanding overparameterized circuit representations via low-rank tensor decompositions.* [🔗](https://openreview.net/forum?id=1btutFdIya)
+        - *Unifying and understanding overparameterized circuit representations via
+          low-rank tensor decompositions.*
+          [🔗](https://openreview.net/forum?id=1btutFdIya)
           Mari, Antonio, Gennaro Vessio, and Antonio Vergari.
           In The 6th Workshop on Tractable Probabilistic Modeling. 2023.