libcirkit 0.1.0__py3-none-any.whl

cirkit/backend/base.py

@@ -0,0 +1,199 @@
+import os
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import IO, Any, Protocol, TypeVar
+
+from cirkit.symbolic.circuit import Circuit
+from cirkit.symbolic.initializers import Initializer
+from cirkit.symbolic.layers import Layer
+from cirkit.symbolic.parameters import ParameterNode
+from cirkit.utils.algorithms import BiMap
+
+CompiledCircuit = TypeVar("CompiledCircuit")
+LayerCompilationSign = type[Layer]
+ParameterCompilationSign = type[ParameterNode]
+InitializerCompilationSign = type[Initializer]
+
+
+class LayerCompilationFunc(Protocol):
+    def __call__(self, compiler: "AbstractCompiler", sl: Layer, **kwargs) -> Any:
+        ...
+
+
+class ParameterCompilationFunc(Protocol):
+    def __call__(self, compiler: "AbstractCompiler", p: ParameterNode, **kwargs) -> Any:
+        ...
+
+
+class InitializerCompilationFunc(Protocol):
+    def __call__(self, compiler: "AbstractCompiler", init: Initializer, **kwargs) -> Any:
+        ...
+
+
+class CompilationRuleNotFound(Exception):
+    def __init__(self, msg: str):
+        super().__init__(msg)
+
+
+SUPPORTED_BACKENDS = ["torch"]
+
+
+class CompiledCircuitsMap:
+    def __init__(self):
+        self._bimap = BiMap[Circuit, CompiledCircuit]()
+
+    def is_compiled(self, sc: Circuit) -> bool:
+        return self._bimap.has_left(sc)
+
+    def has_symbolic(self, cc: CompiledCircuit) -> bool:
+        return self._bimap.has_right(cc)
+
+    def get_compiled_circuit(self, sc: Circuit) -> CompiledCircuit:
+        return self._bimap.get_left(sc)
+
+    def get_symbolic_circuit(self, cc: CompiledCircuit) -> Circuit:
+        return self._bimap.get_right(cc)
+
+    def register_compiled_circuit(self, sc: Circuit, cc: CompiledCircuit):
+        self._bimap.add(sc, cc)
+
+
+class CompilerRegistry:
+    def __init__(
+        self,
+        layer_rules: dict[LayerCompilationSign, LayerCompilationFunc] | None = None,
+        parameter_rules: dict[ParameterCompilationSign, ParameterCompilationFunc] | None = None,
+        initializer_rules: None
+        | (dict[InitializerCompilationSign, InitializerCompilationFunc]) = None,
+    ):
+        self._layer_rules = {} if layer_rules is None else layer_rules
+        self._parameter_rules = {} if parameter_rules is None else parameter_rules
+        self._initializer_rules = {} if initializer_rules is None else initializer_rules
+
+    @staticmethod
+    def _validate_rule_sign(func: Callable, sym_cls: type) -> type | None:
+        args = func.__annotations__
+        if "return" not in args or "compiler" not in args or len(args) != 3:
+            return None
+        if not issubclass(args["compiler"], AbstractCompiler):
+            return None
+        arg_names = list(filter(lambda a: a not in ("return", "compiler"), args.keys()))
+        found_sym_cls = args[arg_names[0]]
+        if not issubclass(found_sym_cls, sym_cls):
+            return None
+        return found_sym_cls
+
+    def add_layer_rule(self, func: LayerCompilationFunc):
+        layer_cls: type[Layer] | None = self._validate_rule_sign(func, Layer)
+        if layer_cls is None:
+            raise ValueError("The function is not a symbolic layer compilation rule")
+        self._layer_rules[layer_cls] = func
+
+    def add_parameter_rule(self, func: ParameterCompilationFunc):
+        param_cls: type[ParameterNode] | None = self._validate_rule_sign(func, ParameterNode)
+        if param_cls is None:
+            raise ValueError("The function is not a symbolic parameter compilation rule")
+        self._parameter_rules[param_cls] = func
+
+    def add_initializer_rule(self, func: InitializerCompilationFunc):
+        init_cls: type[Initializer] | None = self._validate_rule_sign(func, Initializer)
+        if init_cls is None:
+            raise ValueError("The function is not a symbolic initializer compilation rule")
+        self._initializer_rules[init_cls] = func
+
+    def retrieve_layer_rule(self, signature: LayerCompilationSign) -> LayerCompilationFunc:
+        if signature not in self._layer_rules:
+            raise CompilationRuleNotFound(
+                f"Layer compilation rule for signature '{signature}' not found"
+            )
+        return self._layer_rules[signature]
+
+    def retrieve_parameter_rule(
+        self, signature: ParameterCompilationSign
+    ) -> ParameterCompilationFunc:
+        if signature not in self._parameter_rules:
+            raise CompilationRuleNotFound(
+                f"Parameter compilation rule for signature '{signature}' not found"
+            )
+        return self._parameter_rules[signature]
+
+    def retrieve_initializer_rule(
+        self, signature: InitializerCompilationSign
+    ) -> InitializerCompilationFunc:
+        if signature not in self._initializer_rules:
+            raise CompilationRuleNotFound(
+                f"Initializer compilation rule for signature '{signature}' not found"
+            )
+        return self._initializer_rules[signature]
+
+
+class AbstractCompiler(ABC):
+    def __init__(self, registry: CompilerRegistry, **flags):
+        self._registry = registry
+        self._flags = flags
+        self._compiled_circuits = CompiledCircuitsMap()
+
+    def is_compiled(self, sc: Circuit) -> bool:
+        return self._compiled_circuits.is_compiled(sc)
+
+    def has_symbolic(self, cc: CompiledCircuit) -> bool:
+        return self._compiled_circuits.has_symbolic(cc)
+
+    def get_compiled_circuit(self, sc: Circuit) -> CompiledCircuit:
+        return self._compiled_circuits.get_compiled_circuit(sc)
+
+    def get_symbolic_circuit(self, cc: CompiledCircuit) -> Circuit:
+        return self._compiled_circuits.get_symbolic_circuit(cc)
+
+    def register_compiled_circuit(self, sc: Circuit, cc: CompiledCircuit):
+        self._compiled_circuits.register_compiled_circuit(sc, cc)
+
+    def add_layer_rule(self, func: LayerCompilationFunc):
+        self._registry.add_layer_rule(func)
+
+    def add_parameter_rule(self, func: ParameterCompilationFunc):
+        self._registry.add_parameter_rule(func)
+
+    def add_initializer_rule(self, func: InitializerCompilationFunc):
+        self._registry.add_initializer_rule(func)
+
+    def retrieve_layer_rule(self, signature: LayerCompilationSign) -> LayerCompilationFunc:
+        return self._registry.retrieve_layer_rule(signature)
+
+    def retrieve_parameter_rule(
+        self, signature: ParameterCompilationSign
+    ) -> ParameterCompilationFunc:
+        return self._registry.retrieve_parameter_rule(signature)
+
+    def retrieve_initializer_rule(
+        self, signature: InitializerCompilationSign
+    ) -> InitializerCompilationFunc:
+        return self._registry.retrieve_initializer_rule(signature)
+
+    def compile(self, sc: Circuit) -> CompiledCircuit:
+        if self.is_compiled(sc):
+            return self.get_compiled_circuit(sc)
+        return self.compile_pipeline(sc)
+
+    @abstractmethod
+    def compile_layer(self, sl: Layer) -> Any:
+        ...
+
+    @abstractmethod
+    def compile_pipeline(self, sc: Circuit) -> CompiledCircuit:
+        ...
+
+    @abstractmethod
+    def save(
+        self,
+        sym_filepath: IO | os.PathLike | str,
+        compiled_filepath: IO | os.PathLike | str,
+    ):
+        ...
+
+    @staticmethod
+    @abstractmethod
+    def load(
+        sym_filepath: IO | os.PathLike | str, tens_filepath: IO | os.PathLike | str
+    ) -> "AbstractCompiler":
+        ...

cirkit/backend/compiler.py

@@ -0,0 +1,213 @@
+from abc import ABC, abstractmethod
+from typing import Any, Protocol, TypeVar, cast
+
+from cirkit.backend.registry import CompilerRegistry
+from cirkit.symbolic.circuit import Circuit
+from cirkit.symbolic.initializers import Initializer
+from cirkit.symbolic.layers import Layer
+from cirkit.symbolic.parameters import ParameterNode
+from cirkit.utils.algorithms import BiMap
+
+SUPPORTED_BACKENDS = ["torch"]
+
+CompiledCircuit = TypeVar("CompiledCircuit")
+
+
+class CompiledCircuitsMap:
+    def __init__(self):
+        self._bimap = BiMap[Circuit, CompiledCircuit]()
+
+    def is_compiled(self, sc: Circuit) -> bool:
+        return self._bimap.has_left(sc)
+
+    def has_symbolic(self, cc: CompiledCircuit) -> bool:
+        return self._bimap.has_right(cc)
+
+    def get_compiled_circuit(self, sc: Circuit) -> CompiledCircuit:
+        return self._bimap.get_left(sc)
+
+    def get_symbolic_circuit(self, cc: CompiledCircuit) -> Circuit:
+        return self._bimap.get_right(cc)
+
+    def register_compiled_circuit(self, sc: Circuit, cc: CompiledCircuit):
+        self._bimap.add(sc, cc)
+
+
+LayerCompilationSign = type[Layer]
+ParameterCompilationSign = type[ParameterNode]
+InitializerCompilationSign = type[Initializer]
+
+
+class LayerCompilationFunc(Protocol):
+    """The layer compilation function protocol."""
+
+    def __call__(self, compiler: "AbstractCompiler", sl: Layer, **kwargs) -> Any:
+        """Compile a symbolic layer, given a compiler.
+
+        Args:
+            compiler: The compiler.
+            sl: The symbolic layer.
+            **kwargs: The optional arguments for the compilation.
+
+        Returns:
+            A representation of the compiled layer, which depends on the chosen compilation backend.
+        """
+
+
+class ParameterCompilationFunc(Protocol):
+    """The parameter node compilation function protocol."""
+
+    def __call__(self, compiler: "AbstractCompiler", p: ParameterNode, **kwargs) -> Any:
+        """Compile a symbolic parameter node, given a compiler.
+
+        Args:
+            compiler: The compiler.
+            p: The symbolic parameter node.
+            **kwargs: The optional arguments for the compilation.
+
+        Returns:
+            A representation of the compiled parameter node,
+                which depends on the chosen compilation backend.
+        """
+
+
+class InitializerCompilationFunc(Protocol):
+    """The initialization method compilation function protocol."""
+
+    def __call__(self, compiler: "AbstractCompiler", init: Initializer, **kwargs) -> Any:
+        """Compile a symbolic initializer, given a compiler.
+
+        Args:
+            compiler: The compiler.
+            init: The symbolic initializer.
+            **kwargs: The optional arguments for the compilation.
+
+        Returns:
+            A representation of the compiled initializer,
+                which depends on the chosen compilation backend.
+        """
+
+
+class CompilationRuleNotFound(Exception):
+    """An exception that is raised when a compilation rule is not found."""
+
+    def __init__(self, msg: str):
+        """Initializes a compilation rule not found exception.
+
+        Args:
+            msg: The message of the exception.
+        """
+        super().__init__(msg)
+
+
+class CompilerLayerRegistry(CompilerRegistry[LayerCompilationSign, LayerCompilationFunc]):
+    @classmethod
+    def _validate_rule_function(cls, func: LayerCompilationFunc) -> bool:
+        ann = func.__annotations__.copy()
+        del ann["return"]
+        args = tuple(ann.keys())
+        return issubclass(ann[args[-1]], Layer)
+
+    @classmethod
+    def _retrieve_signature(cls, func: LayerCompilationFunc) -> LayerCompilationSign:
+        ann = func.__annotations__.copy()
+        del ann["return"]
+        args = tuple(ann.keys())
+        return cast(LayerCompilationSign, ann[args[-1]])
+
+
+class CompilerParameterRegistry(
+    CompilerRegistry[ParameterCompilationSign, ParameterCompilationFunc]
+):
+    @classmethod
+    def _validate_rule_function(cls, func: ParameterCompilationFunc) -> bool:
+        ann = func.__annotations__.copy()
+        del ann["return"]
+        args = tuple(ann.keys())
+        return issubclass(ann[args[-1]], ParameterNode)
+
+    @classmethod
+    def _retrieve_signature(cls, func: ParameterCompilationFunc) -> ParameterCompilationSign:
+        ann = func.__annotations__.copy()
+        del ann["return"]
+        args = tuple(ann.keys())
+        return cast(ParameterCompilationSign, ann[args[-1]])
+
+
+class CompilerInitializerRegistry(
+    CompilerRegistry[InitializerCompilationSign, InitializerCompilationFunc]
+):
+    @classmethod
+    def _validate_rule_function(cls, func: InitializerCompilationFunc) -> bool:
+        ann = func.__annotations__.copy()
+        del ann["return"]
+        args = tuple(ann.keys())
+        return issubclass(ann[args[-1]], Initializer)
+
+    @classmethod
+    def _retrieve_signature(cls, func: ParameterCompilationFunc) -> InitializerCompilationSign:
+        ann = func.__annotations__.copy()
+        del ann["return"]
+        args = tuple(ann.keys())
+        return cast(InitializerCompilationSign, ann[args[-1]])
+
+
+class AbstractCompiler(ABC):
+    def __init__(
+        self,
+        layers_registry: CompilerLayerRegistry,
+        parameters_registry: CompilerParameterRegistry,
+        initializers_registry: CompilerInitializerRegistry,
+        **flags,
+    ):
+        self._layers_registry = layers_registry
+        self._parameters_registry = parameters_registry
+        self._initializers_registry = initializers_registry
+        self._flags = flags
+        self._compiled_circuits = CompiledCircuitsMap()
+
+    def is_compiled(self, sc: Circuit) -> bool:
+        return self._compiled_circuits.is_compiled(sc)
+
+    def has_symbolic(self, cc: CompiledCircuit) -> bool:
+        return self._compiled_circuits.has_symbolic(cc)
+
+    def get_compiled_circuit(self, sc: Circuit) -> CompiledCircuit:
+        return self._compiled_circuits.get_compiled_circuit(sc)
+
+    def get_symbolic_circuit(self, cc: CompiledCircuit) -> Circuit:
+        return self._compiled_circuits.get_symbolic_circuit(cc)
+
+    def register_compiled_circuit(self, sc: Circuit, cc: CompiledCircuit):
+        self._compiled_circuits.register_compiled_circuit(sc, cc)
+
+    def add_layer_rule(self, func: LayerCompilationFunc):
+        self._layers_registry.add_rule(func)
+
+    def add_parameter_rule(self, func: ParameterCompilationFunc):
+        self._parameters_registry.add_rule(func)
+
+    def add_initializer_rule(self, func: InitializerCompilationFunc):
+        self._initializers_registry.add_rule(func)
+
+    def retrieve_layer_rule(self, signature: LayerCompilationSign) -> LayerCompilationFunc:
+        return self._layers_registry.retrieve_rule(signature)
+
+    def retrieve_parameter_rule(
+        self, signature: ParameterCompilationSign
+    ) -> ParameterCompilationFunc:
+        return self._parameters_registry.retrieve_rule(signature)
+
+    def retrieve_initializer_rule(
+        self, signature: InitializerCompilationSign
+    ) -> InitializerCompilationFunc:
+        return self._initializers_registry.retrieve_rule(signature)
+
+    def compile(self, sc: Circuit) -> CompiledCircuit:
+        if self.is_compiled(sc):
+            return self.get_compiled_circuit(sc)
+        return self.compile_pipeline(sc)
+
+    @abstractmethod
+    def compile_pipeline(self, sc: Circuit) -> CompiledCircuit:
+        ...

cirkit/backend/registry.py

@@ -0,0 +1,53 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+RegistrySign = TypeVar("RegistrySign")
+RegistryFunc = TypeVar("RegistryFunc")
+
+
+class InvalidRuleSign(Exception):
+    def __init__(self, annotations: dict[str, type]):
+        super().__init__(
+            f"Cannot extract rule signature from function with annotations '{annotations}"
+        )
+
+
+class InvalidRuleFunction(Exception):
+    def __init__(self, annotations: dict[str, type]):
+        super().__init__(f"Invalid Compilation rule function with annotations '{annotations}'")
+
+
+class CompilationRuleNotFound(Exception):
+    def __init__(self, signature: RegistrySign):
+        super().__init__(f"Compilation rule for signature '{signature}' not found")
+
+
+class CompilerRegistry(Generic[RegistrySign, RegistryFunc], ABC):
+    def __init__(self, rules: dict[RegistrySign, RegistryFunc] | None = None):
+        self._rules = {} if rules is None else rules
+
+    @classmethod
+    @abstractmethod
+    def _validate_rule_function(cls, func: RegistryFunc) -> bool:
+        ...
+
+    @classmethod
+    def _retrieve_signature(cls, func: RegistryFunc) -> RegistrySign:
+        raise InvalidRuleSign(func.__annotations__)
+
+    @property
+    def signatures(self) -> list[RegistrySign]:
+        return list(self._rules)
+
+    def add_rule(self, func: RegistryFunc, *, signature: RegistrySign | None = None) -> None:
+        if not self._validate_rule_function(func):
+            raise InvalidRuleFunction(func.__annotations__)
+        if signature is None:
+            signature = self._retrieve_signature(func)
+        self._rules[signature] = func
+
+    def retrieve_rule(self, signature: RegistrySign) -> RegistryFunc:
+        func = self._rules.get(signature, None)
+        if func is not None:
+            return func
+        raise CompilationRuleNotFound(signature)

cirkit/backend/torch/circuits.py

@@ -0,0 +1,217 @@
+from collections.abc import Callable, Iterator, Sequence
+
+import torch
+from torch import Tensor
+
+from cirkit.backend.torch.graph.folding import (
+    build_address_book_stacked_entry,
+    build_unfold_index_info,
+)
+from cirkit.backend.torch.graph.modules import (
+    AddressBook,
+    AddressBookEntry,
+    FoldIndexInfo,
+    TorchDiAcyclicGraph,
+)
+from cirkit.backend.torch.layers import TorchInputLayer, TorchLayer
+from cirkit.symbolic.circuit import StructuralProperties
+from cirkit.utils.scope import Scope
+
+
+class LayerAddressBook(AddressBook):
+    def __init__(self, entries: list[AddressBookEntry]):
+        super().__init__(entries)
+
+    def lookup(
+        self, module_outputs: list[Tensor], *, in_graph: Tensor | None = None
+    ) -> Iterator[tuple[TorchLayer | None, tuple[Tensor, ...]]]:
+        # Loop through the entries and yield inputs
+        for entry in self._entries:
+            # Catch the case there are some inputs coming from other modules
+            if entry.in_module_ids:
+                (in_fold_idx,) = entry.in_fold_idx
+                (in_module_ids,) = entry.in_module_ids
+                if len(in_module_ids) == 1:
+                    x = module_outputs[in_module_ids[0]]
+                else:
+                    x = torch.cat([module_outputs[mid] for mid in in_module_ids], dim=0)
+                x = x[in_fold_idx]
+                yield entry.module, (x,)
+                continue
+
+            # Catch the case there are no inputs coming from other modules
+            # That is, we are gathering the inputs of input layers
+            assert isinstance(entry.module, TorchInputLayer)
+            if in_graph is None:
+                yield entry.module, ()
+            else:
+                # in_graph: An input batch (assignments to variables) of shape (B, C, D)
+                # scope_idx: The scope of the layers in each fold, a tensor of shape (F, D'), D' < D
+                # x: (B, C, D) -> (B, C, F, D') -> (F, C, B, D')
+                x = in_graph[..., entry.module.scope_idx].permute(2, 1, 0, 3)
+                yield entry.module, (x,)
+
+    @classmethod
+    def from_index_info(
+        cls,
+        fold_idx_info: FoldIndexInfo,
+        *,
+        incomings_fn: Callable[[TorchLayer], Sequence[TorchLayer]],
+    ) -> "LayerAddressBook":
+        # The address book entries being built
+        entries: list[AddressBookEntry] = []
+
+        # A useful dictionary mapping module ids to their number of folds
+        num_folds: dict[int, int] = {}
+
+        # Build the bookkeeping data structure by following the topological ordering
+        for mid, m in enumerate(fold_idx_info.ordering):
+            # Retrieve the index information of the input modules
+            in_modules_fold_idx = fold_idx_info.in_fold_idx[mid]
+
+            # Catch the case of a folded module having the output of another module as input
+            if incomings_fn(m):
+                entry = build_address_book_stacked_entry(
+                    m, in_modules_fold_idx, num_folds=num_folds
+                )
+            else:
+                # Catch the case of a folded module having the input of the network as input
+                # That is, this is the case of an input layer
+                entry = AddressBookEntry(m, [], [])
+
+            num_folds[mid] = m.num_folds
+            entries.append(entry)
+
+        # Append the last bookkeeping entry with the information to compute the output tensor
+        entry = build_address_book_stacked_entry(
+            None, [fold_idx_info.out_fold_idx], num_folds=num_folds, output=True
+        )
+        entries.append(entry)
+
+        return LayerAddressBook(entries)
+
+
+class AbstractTorchCircuit(TorchDiAcyclicGraph[TorchLayer]):
+    def __init__(
+        self,
+        scope: Scope,
+        num_channels: int,
+        layers: Sequence[TorchLayer],
+        in_layers: dict[TorchLayer, Sequence[TorchLayer]],
+        outputs: Sequence[TorchLayer],
+        *,
+        properties: StructuralProperties,
+        fold_idx_info: FoldIndexInfo | None = None,
+    ) -> None:
+        super().__init__(
+            layers,
+            in_layers,
+            outputs,
+            fold_idx_info=fold_idx_info,
+        )
+        self._scope = scope
+        self._num_channels = num_channels
+        self._properties = properties
+
+    @property
+    def scope(self) -> Scope:
+        return self._scope
+
+    @property
+    def num_variables(self) -> int:
+        return len(self.scope)
+
+    @property
+    def num_channels(self) -> int:
+        return self._num_channels
+
+    @property
+    def properties(self) -> StructuralProperties:
+        return self._properties
+
+    @property
+    def layers(self) -> Sequence[TorchLayer]:
+        return self.nodes
+
+    def layer_inputs(self, l: TorchLayer) -> Sequence[TorchLayer]:
+        return self.node_inputs(l)
+
+    def layer_outputs(self, l: TorchLayer) -> Sequence[TorchLayer]:
+        return self.node_outputs(l)
+
+    @property
+    def layers_inputs(self) -> dict[TorchLayer, Sequence[TorchLayer]]:
+        return self.nodes_inputs
+
+    @property
+    def layers_outputs(self) -> dict[TorchLayer, Sequence[TorchLayer]]:
+        return self.nodes_outputs
+
+    def reset_parameters(self) -> None:
+        # For each layer, initialize its parameters, if any
+        for l in self.layers:
+            for p in l.params.values():
+                p.reset_parameters()
+
+    def _set_device(self, device: str | torch.device | int) -> None:
+        for l in self.layers:
+            for p in l.params.values():
+                p._set_device(device)
+        super()._set_device(device)
+
+    def _build_unfold_index_info(self) -> FoldIndexInfo:
+        return build_unfold_index_info(
+            self.topological_ordering(), outputs=self.outputs, incomings_fn=self.node_inputs
+        )
+
+    def _build_address_book(self, fold_idx_info: FoldIndexInfo) -> LayerAddressBook:
+        return LayerAddressBook.from_index_info(fold_idx_info, incomings_fn=self.layer_inputs)
+
+    def _evaluate_layers(self, x: Tensor) -> Tensor:
+        # Evaluate layers on the given input
+        y = self.evaluate(x)  # (O, B, K)
+        return y.transpose(0, 1)  # (B, O, K)
+
+
+class TorchCircuit(AbstractTorchCircuit):
+    """The tensorized circuit with concrete computational graph in PyTorch.
+
+    This class is aimed for computation, and therefore does not include structural properties.
+    """
+
+    def __call__(self, x: Tensor) -> Tensor:
+        """Invoke the forward function.
+
+        Args:
+            x (Tensor): The input of the circuit, shape (B, C, D).
+
+        Returns:
+            Tensor: The output of the circuit, shape (B, num_out, num_cls).
+        """  # TODO: single letter name?
+        # IGNORE: Idiom for nn.Module.__call__.
+        return super().__call__(x)  # type: ignore[no-any-return,misc]
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._evaluate_layers(x)
+
+
+class TorchConstantCircuit(AbstractTorchCircuit):
+    """The tensorized circuit with concrete computational graph in PyTorch.
+
+    This class is aimed for computation, and therefore does not include strutural properties.
+    """
+
+    def __call__(self) -> Tensor:
+        """Invoke the forward function.
+
+        Returns:
+            Tensor: The output of the circuit, shape (B, num_out, num_cls).
+        """  # TODO: single letter name?
+        # IGNORE: Idiom for nn.Module.__call__.
+        return super().__call__()  # type: ignore[no-any-return,misc]
+
+    def forward(self) -> Tensor:
+        # Evaluate the layers using some dummy input
+        x = torch.empty(size=(1, self.num_channels, self.num_variables), device=self.device)
+        x = self._evaluate_layers(x)  # (B, O, K)
+        return x.squeeze(dim=0)  # (O, K)