diff --git a/cirkit/backend/torch/compiler.py b/cirkit/backend/torch/compiler.py index 634bdbe1..26a143ff 100644 --- a/cirkit/backend/torch/compiler.py +++ b/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. diff --git a/cirkit/backend/torch/graph/optimize.py b/cirkit/backend/torch/graph/optimize.py index d795d6f3..be2fbab9 100644 --- a/cirkit/backend/torch/graph/optimize.py +++ b/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:] diff --git a/cirkit/backend/torch/optimization/layers.py b/cirkit/backend/torch/optimization/layers.py index e6f96290..c0b9824d 100644 --- a/cirkit/backend/torch/optimization/layers.py +++ b/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, diff --git a/cirkit/backend/torch/optimization/parameters.py b/cirkit/backend/torch/optimization/parameters.py index 1676cb98..c58595d2 100644 --- a/cirkit/backend/torch/optimization/parameters.py +++ b/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 diff --git a/cirkit/backend/torch/parameters/nodes.py b/cirkit/backend/torch/parameters/nodes.py index e98a6b7f..8a7ce0d7 100644 --- a/cirkit/backend/torch/parameters/nodes.py +++ b/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): diff --git a/cirkit/backend/torch/parameters/pic.py b/cirkit/backend/torch/parameters/pic.py index a6ea4b6b..a09449d9 100644 --- a/cirkit/backend/torch/parameters/pic.py +++ b/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, diff --git a/cirkit/backend/torch/semiring.py b/cirkit/backend/torch/semiring.py index bd1958d2..eb54f7ac 100644 --- a/cirkit/backend/torch/semiring.py +++ b/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: ... diff --git a/cirkit/backend/torch/utils.py b/cirkit/backend/torch/utils.py index f961d386..77ea8b1a 100644 --- a/cirkit/backend/torch/utils.py +++ b/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) diff --git a/cirkit/symbolic/io.py b/cirkit/symbolic/io.py index 86121be0..c055f63f 100644 --- a/cirkit/symbolic/io.py +++ b/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", diff --git a/cirkit/templates/logic/graph.py b/cirkit/templates/logic/graph.py index 8a7d8095..17ab91cc 100644 --- a/cirkit/templates/logic/graph.py +++ b/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: diff --git a/cirkit/templates/logic/sdd.py b/cirkit/templates/logic/sdd.py index 6f5f3472..3e044b62 100644 --- a/cirkit/templates/logic/sdd.py +++ b/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: diff --git a/cirkit/templates/region_graph/algorithms/chow_liu.py b/cirkit/templates/region_graph/algorithms/chow_liu.py index 47110609..1c50a681 100644 --- a/cirkit/templates/region_graph/algorithms/chow_liu.py +++ b/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 diff --git a/cirkit/templates/region_graph/algorithms/quad.py b/cirkit/templates/region_graph/algorithms/quad.py index 39c5ec68..0e3ee160 100644 --- a/cirkit/templates/region_graph/algorithms/quad.py +++ b/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. diff --git a/cirkit/templates/region_graph/algorithms/random.py b/cirkit/templates/region_graph/algorithms/random.py index a392104f..b5b01bde 100644 --- a/cirkit/templates/region_graph/algorithms/random.py +++ b/cirkit/templates/region_graph/algorithms/random.py @@ -20,8 +20,11 @@ def RandomBinaryTree( """Construct a RG with random binary trees. See: - - *Random sum-product networks: A simple but effective approach to probabilistic deep learning.* [πŸ”—](https://proceedings.mlr.press/v115/peharz20a.html) - Robert Peharz, Antonio Vergari, Karl Stelzner, Alejandro Molina, Xiaoting Shao, Martin Trapp, Kristian Kersting, and Zoubin Ghahramani. + - *Random sum-product networks: A simple but effective approach to probabilistic + deep learning.* + [πŸ”—](https://proceedings.mlr.press/v115/peharz20a.html) + Robert Peharz, Antonio Vergari, Karl Stelzner, Alejandro Molina, Xiaoting Shao, + Martin Trapp, Kristian Kersting, and Zoubin Ghahramani. In Uncertainty in Artificial Intelligence, pp. 334-344. PMLR, 2020. Args: diff --git a/cirkit/templates/region_graph/graph.py b/cirkit/templates/region_graph/graph.py index b481bedf..fb7c83a9 100644 --- a/cirkit/templates/region_graph/graph.py +++ b/cirkit/templates/region_graph/graph.py @@ -341,7 +341,7 @@ def load(filename: str) -> "RegionGraph": return RegionGraph(nodes, in_nodes, outputs=outputs) # TODO: refactor the following method as to simplify its structure (e.g., remove inline methods) - def build_circuit( + def build_circuit( # pylint: disable=too-many-locals,too-many-branches,too-many-statements self, *, input_factory: InputLayerFactory | Mapping[Scope, InputLayerFactory], @@ -365,7 +365,8 @@ def build_circuit( Args: input_factory: A factory that builds an input layer. If a Mapping between scopes and - input layer factories is given, then a different input layer is built for each feature. + input layer factories is given, then a different input layer is built for + each feature. In this case, an input layer should be specified for each variable/scope. sum_product: The sum-product layer to use. It can be None, 'cp', 'cp-t', or 'tucker'. sum_weight_factory: The factory to construct the weights of the sum layers. diff --git a/cirkit/templates/tensor_factorizations.py b/cirkit/templates/tensor_factorizations.py index c8ea31df..740d9112 100644 --- a/cirkit/templates/tensor_factorizations.py +++ b/cirkit/templates/tensor_factorizations.py @@ -59,7 +59,8 @@ def cp( with additional weights, i.e., a CP factorization of the form $$ - c(X_1,\ldots,X_n) = t_{X_1\cdots X_n} = \sum_{i=1}^R w_i \: a^{(1)}_{X_1 i} \ldots a^{(n)}_{X_n i}, + c(X_1,\ldots,X_n) = t_{X_1\cdots X_n} + = \sum_{i=1}^R w_i \: a^{(1)}_{X_1 i} \ldots a^{(n)}_{X_n i}, $$ where $\mathbf{w}\in\mathbb{R}^R$ are additional weights. @@ -74,7 +75,8 @@ def cp( $\mathcal{T}\in\mathbb{R}_+^{I_1\times \ldots\times I_n}$ as the distribution $$ - p(X_1,\ldots,X_n) = t_{X_1\cdots X_n} = \sum_{i=1}^R p(Z=i) \: p(X_1\mid Z=i) \cdots p(X_n\mid Z=i), + p(X_1,\ldots,X_n) = t_{X_1\cdots X_n} + = \sum_{i=1}^R p(Z=i) \: p(X_1\mid Z=i) \cdots p(X_n\mid Z=i), $$ where $Z$ is a discrete latent variable modelled by $p(Z)$. @@ -150,12 +152,14 @@ def tucker( and $c$ computes a rank-$R$ Tucker factorization, i.e., $$ - c(X_1,\ldots,X_n) = t_{X_1\cdots X_n} = \sum_{r_1=1}^R \cdots \sum_{r_n=1}^R w_{r_1\cdots r_n} a^{(1)}_{X_1 r_1} \cdots a^{(n)}_{X_n r_n}, + c(X_1,\ldots,X_n) = t_{X_1\cdots X_n} + = \sum_{r_1=1}^R \cdots \sum_{r_n=1}^R w_{r_1\cdots r_n} + a^{(1)}_{X_1 r_1} \cdots a^{(n)}_{X_n r_n}, $$ - where for $1\leq j\leq n$ we have that $\mathbf{A}^{(j)}\in\mathbb{R}^{I_j\times R}$ is the $j$-th factor, - and $\mathcal{W}\in\mathbb{R}^{R\times\cdots\times R}$ is an $n$-dimensional tensor, sometimes called - the core tensor of the Tucker factorization. + where for $1\leq j\leq n$ we have that $\mathbf{A}^{(j)}\in\mathbb{R}^{I_j\times R}$ is the + $j$-th factor, and $\mathcal{W}\in\mathbb{R}^{R\times\cdots\times R}$ is an $n$-dimensional + tensor, sometimes called the core tensor of the Tucker factorization. This method allows you to specify different types of parameterizations for the factors and the core tensor. For example, if the arguments ```factor_param``` and @@ -167,7 +171,9 @@ def tucker( $\mathcal{T}\in\mathbb{R}_+^{I_1\times \ldots\times I_n}$ as the distribution $$ - p(X_1,\ldots,X_n) = t_{X_1\cdots X_n} = \sum_{r_1=1}^R \cdots \sum_{r_n=1}^R p(Z=(r_1,\ldots,r_n)) \: p(X_1\mid Z=r_1) \cdots p(X_n\mid Z=r_n), + p(X_1,\ldots,X_n) = t_{X_1\cdots X_n} + = \sum_{r_1=1}^R \cdots \sum_{r_n=1}^R p(Z=(r_1,\ldots,r_n)) \: p(X_1\mid Z=r_1) + \cdots p(X_n\mid Z=r_n), $$ where $Z$ is a discrete latent variable taking value in $\{1,\ldots,R\}^n$ and modelled by @@ -244,7 +250,9 @@ def tensor_train( and $c$ computes a rank-$R$ TT/MPS factorization, i.e., $$ - c(X_1,\ldots,X_n) = t_{X_1\cdots X_n} = \sum_{r_1=1}^R \cdots \sum_{r_{n-1}=1}^R v^{(1)}_{X_1 r_1} v^{(2)}_{X_2 r_1 r_2} \cdots v^{(n-1)}_{X_{n-1} r_{n-2} r_{n-1}} v^{(n)}_{X_n r_{n-1}}, pylint: disable=line-too-long + c(X_1,\ldots,X_n) = t_{X_1\cdots X_n} + = \sum_{r_1=1}^R \cdots \sum_{r_{n-1}=1}^R v^{(1)}_{X_1 r_1} v^{(2)}_{X_2 r_1 r_2} + \cdots v^{(n-1)}_{X_{n-1} r_{n-2} r_{n-1}} v^{(n)}_{X_n r_{n-1}}, $$ where $\mathbf{V}^{(1)}\in\mathbb{R}^{I_1\times R}$,