Merge pull request #168 from april-tools/fix_pipeline

Fix pipeline and add test for it

cirkit/new/__init__.py

@@ -1 +1,29 @@
-# TODO: we need to decide what to import here
+from typing import Type, Union, cast
+
+from . import layers as layers
+from . import model as model
+from . import region_graph as region_graph
+from . import reparams as reparams
+from . import symbolic as symbolic
+from . import utils as utils
+
+
+def set_layer_comp_space(comp_space: Union[Type[utils.ComputationSapce], str]) -> None:
+    """Set the global computational space for layers.
+
+    Args:
+        comp_space (Union[Type[utils.ComputationSapce], str]): The computational space to use, can \
+            be speficied by either the class or the name.
+    """
+    if isinstance(comp_space, str):
+        comp_space = utils.ComputationSapce.get_comp_space_by_name(comp_space)
+
+    # Cast: a cast is needed because __final__ may be undefined.
+    assert cast(
+        bool, getattr(comp_space, "__final__", False)
+    ), "A usable ComputationSapce must be final."
+
+    layers.Layer.comp_space = comp_space
+
+
+set_layer_comp_space(utils.LogSpace)

cirkit/new/layers/layer.py

@@ -5,13 +5,13 @@
 from torch import Tensor, nn
 
 from cirkit.new.reparams import Reparameterization
-from cirkit.new.utils import ComputationSapce, LogSpace
+from cirkit.new.utils import ComputationSapce
 
 
 class Layer(nn.Module, ABC):
     """The abstract base class for all layers."""
 
-    comp_space: Type[ComputationSapce] = LogSpace
+    comp_space: Type[ComputationSapce]
 
     # Disable: reparam is not used in the base class. It's only here for the interface.
     def __init__(

cirkit/new/region_graph/algorithms/poon_domingos.py

@@ -149,7 +149,8 @@ def PoonDomingos(
     queue: Deque[HyperCube] = deque()
     depth_dict: Dict[HyperCube, int] = {}  # Also serve as a "visited" set.
 
-    # A cast is required to get rid of Literal[0].  # TODO: type checking bug? annotation not work
+    # TODO: type checking bug? annotation not work
+    # Cast: A cast is required to get rid of Literal[0].
     cur_hypercube = cast(HyperCube, ((0,) * len(shape), tuple(shape)))
     graph.add_node(RegionNode(hypercube_to_scope[cur_hypercube]))
     queue.append(cur_hypercube)

cirkit/new/region_graph/rg_node.py

@@ -32,7 +32,7 @@ def __repr__(self) -> str:
         Returns:
             str: The str representation of the node.
         """
-        return f"{self.__class__.__name__}({self.scope})"
+        return f"{self.__class__.__name__}@0x{id(self):x}({self.scope})"
 
     # __hash__ and __eq__ are defined by default to compare on object identity, i.e.,
     # (a is b) <=> (a == b) <=> (hash(a) == hash(b)).

cirkit/new/reparams/reparam.py

@@ -83,7 +83,7 @@ def materialize(
                 Defaults to None.
         """
         # NOTE: Subclasses should never call into this materialize() when is_materialized.
-        assert self.is_materialized, "This reparameterization is already materialized."
+        assert not self.is_materialized, "This reparameterization is already materialized."
 
         self.shape = tuple(shape)
 

cirkit/new/symbolic/symbolic_circuit.py

@@ -1,6 +1,13 @@
-from typing import Any, Dict, Iterable, Iterator, Optional, Type, Union
+from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Type, Union
 
-from cirkit.new.layers import InputLayer, MixingLayer, ProductLayer, SumLayer, SumProductLayer
+from cirkit.new.layers import (
+    DenseLayer,
+    InputLayer,
+    MixingLayer,
+    ProductLayer,
+    SumLayer,
+    SumProductLayer,
+)
 from cirkit.new.region_graph import PartitionNode, RegionGraph, RegionNode, RGNode
 from cirkit.new.reparams import Reparameterization
 from cirkit.new.symbolic.symbolic_layer import (
@@ -11,13 +18,14 @@
 )
 from cirkit.new.utils import OrderedSet, Scope
 
-# TODO: double check docs and __repr__
+# TODO: __repr__?
 
 
 # Disable: It's designed to have these many attributes.
 class SymbolicTensorizedCircuit:  # pylint: disable=too-many-instance-attributes
     """The symbolic representation of a tensorized circuit."""
 
+    # TODO: is this the best way to provide reparam? or give a layer-wise mapping?
     # TODO: how to design interface? require kwargs only?
     # TODO: how to deal with too-many?
     # pylint: disable-next=too-many-arguments,too-many-locals
@@ -30,10 +38,10 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
         num_classes: int = 1,
         input_layer_cls: Type[InputLayer],
         input_layer_kwargs: Optional[Dict[str, Any]] = None,
-        input_reparam: Optional[Reparameterization] = None,
+        input_reparam: Callable[[], Optional[Reparameterization]] = lambda: None,
         sum_layer_cls: Type[Union[SumLayer, SumProductLayer]],
         sum_layer_kwargs: Optional[Dict[str, Any]] = None,
-        sum_reparam: Reparameterization,
+        sum_reparam: Callable[[], Reparameterization],
         prod_layer_cls: Type[Union[ProductLayer, SumProductLayer]],
         prod_layer_kwargs: Optional[Dict[str, Any]] = None,
     ):
@@ -49,14 +57,16 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
             input_layer_cls (Type[InputLayer]): The layer class for input layers.
             input_layer_kwargs (Optional[Dict[str, Any]], optional): The additional kwargs for \
                 input layer class. Defaults to None.
-            input_reparam (Optional[Reparameterization], optional): The reparameterization for \
-                input layer parameters, can be None if it has no params. Defaults to None.
+            input_reparam (Callable[[], Optional[Reparameterization]], optional): The factory to \
+                construct reparameterizations for input layer parameters, can produce None if no \
+                params is needed. Defaults to lambda: None.
             sum_layer_cls (Type[Union[SumLayer, SumProductLayer]]): The layer class for sum \
                 layers, can be either just a class of SumLayer, or a class of SumProductLayer to \
-                indicate layer fusion..
+                indicate layer fusion.
             sum_layer_kwargs (Optional[Dict[str, Any]], optional): The additional kwargs for sum \
                 layer class. Defaults to None.
-            sum_reparam (Reparameterization): The reparameterization for sum layer parameters.
+            sum_reparam (Callable[[], Reparameterization]): The factory to construct \
+                reparameterizations for sum layer parameters.
             prod_layer_cls (Type[Union[ProductLayer, SumProductLayer]]): The layer class for \
                 product layers, can be either just a class of ProductLayer, or a class of \
                 SumProductLayer to indicate layer fusion.
@@ -91,17 +101,17 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
                         num_units=num_input_units,
                         layer_cls=input_layer_cls,
                         layer_kwargs=input_layer_kwargs,  # type: ignore[misc]
-                        reparam=input_reparam,
+                        reparam=input_reparam(),
                     )
                 ]
                 # This also works when the input is also output, in which case num_classes is used.
                 layer_out = SymbolicSumLayer(
                     rg_node,
                     layers_in,
                     num_units=num_sum_units if rg_node.outputs else num_classes,
-                    layer_cls=sum_layer_cls,
-                    layer_kwargs=sum_layer_kwargs,  # type: ignore[misc]
-                    reparam=sum_reparam,
+                    layer_cls=DenseLayer,  # TODO: can be other sum layer, but how to pass in???
+                    layer_kwargs={},  # type: ignore[misc]
+                    reparam=sum_reparam(),
                 )
             elif isinstance(rg_node, RegionNode) and len(rg_node.inputs) == 1:  # Simple inner.
                 # layers_in keeps the same.
@@ -111,7 +121,7 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
                     num_units=num_sum_units if rg_node.outputs else num_classes,
                     layer_cls=sum_layer_cls,
                     layer_kwargs=sum_layer_kwargs,  # type: ignore[misc]
-                    reparam=sum_reparam,
+                    reparam=sum_reparam(),
                 )
             elif isinstance(rg_node, RegionNode) and len(rg_node.inputs) > 1:  # Inner with mixture.
                 # MixingLayer cannot change number of units, so must project early.
@@ -122,7 +132,7 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
                         num_units=num_sum_units if rg_node.outputs else num_classes,
                         layer_cls=sum_layer_cls,
                         layer_kwargs=sum_layer_kwargs,  # type: ignore[misc]
-                        reparam=sum_reparam,
+                        reparam=sum_reparam(),
                     )
                     for layer_in in layers_in
                 ]
@@ -132,7 +142,7 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
                     num_units=num_sum_units if rg_node.outputs else num_classes,
                     layer_cls=MixingLayer,
                     layer_kwargs={},  # type: ignore[misc]
-                    reparam=sum_reparam,  # TODO: use a constant reparam here?
+                    reparam=sum_reparam(),  # TODO: use a constant reparam here?
                 )
             elif isinstance(rg_node, PartitionNode):
                 # layers_in keeps the same.

cirkit/new/symbolic/symbolic_layer.py

@@ -5,8 +5,6 @@
 from cirkit.new.region_graph import PartitionNode, RegionNode, RGNode
 from cirkit.new.reparams import Reparameterization
 
-# TODO: double check __repr__
-
 
 # Disable: It's intended for SymbolicLayer to have these many attrs.
 class SymbolicLayer(ABC):  # pylint: disable=too-many-instance-attributes
@@ -56,15 +54,20 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
         self.layer_kwargs = layer_kwargs if layer_kwargs is not None else {}  # type: ignore[misc]
         self.reparam = reparam
 
-    # We require subclasses to implement __repr__ on their own. This also forbids the instantiation
-    # of this abstract class.
-    @abstractmethod
     def __repr__(self) -> str:
         """Generate the repr string of the layer.
 
         Returns:
             str: The str representation of the layer.
         """
+        repr_kv = ", ".join(f"{k}={v}" for k, v in self._repr_dict.items())
+        return f"{self.__class__.__name__}@0x{id(self):x}({repr_kv})"
+
+    # We use an abstract instead of direct attribute so that this class includes an abstract method.
+    @property
+    @abstractmethod
+    def _repr_dict(self) -> Dict[str, object]:  # Use object to avoid Any.
+        """The dict of key-value pairs used in __repr__."""
 
     # __hash__ and __eq__ are defined by default to compare on object identity, i.e.,
     # (a is b) <=> (a == b) <=> (hash(a) == hash(b)).
@@ -129,22 +132,18 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
         )
         assert self.inputs, "SymbolicSumLayer must be an inner layer of the SymbC."
 
-    def __repr__(self) -> str:
-        """Generate the repr string of the layer.
-
-        Returns:
-            str: The str representation of the layer.
-        """
-        class_name = self.__class__.__name__
-        layer_cls_name = self.layer_cls.__name__
-
-        return (
-            f"{class_name}:\n"  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
-            f"Scope: {repr(self.scope)}\n"
-            f"Layer Class: {layer_cls_name}\n"
-            f"Layer KWArgs: {repr(self.layer_kwargs)}\n"
-            f"Number of Units: {repr(self.num_units)}\n"
-        )
+    @property
+    def _repr_dict(self) -> Dict[str, object]:
+        """The dict of key-value pairs used in __repr__."""
+        # Ignore: Unavoidable for kwargs.
+        return {
+            "scope": self.scope,
+            "arity": self.arity,
+            "num_units": self.num_units,
+            "layer_cls": self.layer_cls,
+            "layer_kwargs": self.layer_kwargs,  # type: ignore[misc]
+            "reparam": self.reparam,  # TODO: repr of reparam
+        }
 
 
 # Disable: It's intended for SymbolicProductLayer to have only these methods.
@@ -188,21 +187,17 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
             reparam=None,
         )
 
-    def __repr__(self) -> str:
-        """Generate the repr string of the layer.
-
-        Returns:
-            str: The str representation of the layer.
-        """
-        class_name = self.__class__.__name__
-        layer_cls_name = self.layer_cls.__name__
-
-        return (
-            f"{class_name}:\n"
-            f"Scope: {repr(self.scope)}\n"
-            f"Layer Class: {layer_cls_name}\n"
-            f"Number of Units: {repr(self.num_units)}\n"
-        )
+    @property
+    def _repr_dict(self) -> Dict[str, object]:
+        """The dict of key-value pairs used in __repr__."""
+        # Ignore: Unavoidable for kwargs.
+        return {
+            "scope": self.scope,
+            "arity": self.arity,
+            "num_units": self.num_units,
+            "layer_cls": self.layer_cls,
+            "layer_kwargs": self.layer_kwargs,  # type: ignore[misc]
+        }
 
 
 # Disable: It's intended for SymbolicInputLayer to have only these methods.
@@ -243,19 +238,14 @@ def __init__(  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
         )
         assert not self.inputs, "SymbolicInputLayer must be an input layer of the SymbC."
 
-    def __repr__(self) -> str:
-        """Generate the repr string of the layer.
-
-        Returns:
-            str: The str representation of the layer.
-        """
-        class_name = self.__class__.__name__
-        layer_cls_name = self.layer_cls.__name__ if self.layer_cls else "None"
-
-        return (
-            f"{class_name}:\n"  # type: ignore[misc]  # Ignore: Unavoidable for kwargs.
-            f"Scope: {repr(self.scope)}\n"
-            f"Input Exp Family Class: {layer_cls_name}\n"
-            f"Layer KWArgs: {repr(self.layer_kwargs)}\n"
-            f"Number of Units: {repr(self.num_units)}\n"
-        )
+    @property
+    def _repr_dict(self) -> Dict[str, object]:
+        """The dict of key-value pairs used in __repr__."""
+        # Ignore: Unavoidable for kwargs.
+        return {
+            "scope": self.scope,
+            "num_units": self.num_units,
+            "layer_cls": self.layer_cls,
+            "layer_kwargs": self.layer_kwargs,  # type: ignore[misc]
+            "reparam": self.reparam,  # TODO: repr of reparam
+        }

cirkit/new/utils/comp_space/comp_space.py

@@ -1,11 +1,31 @@
 from abc import ABC, abstractmethod
-from typing import Callable, Optional, Sequence, Union
-from typing_extensions import TypeVarTuple, Unpack  # TODO: in typing from 3.11
+from typing import Callable, Dict, Optional, Sequence, Type, TypeVar, Union, cast
+from typing_extensions import Self, TypeVarTuple, Unpack  # TODO: in typing from 3.11
 
 from torch import Tensor
 
 Ts = TypeVarTuple("Ts")
 
+CompSpaceT = TypeVar("CompSpaceT", bound=Type["ComputationSapce"])
+
+
+def register_comp_space(cls: CompSpaceT) -> CompSpaceT:
+    """Register a concrete ComputationSapce by its name.
+
+    Args:
+        cls (CompSpaceT): The ComputationSapce subclass to register.
+
+    Returns:
+        CompSpaceT: The class passed in.
+    """
+    # Cast: a cast is needed because __final__ may be undefined.
+    assert cast(
+        bool, getattr(cls, "__final__", False)
+    ), "Subclasses of ComputationSapce should be final."
+    # Disable: Access to _registry is by design.
+    ComputationSapce._registry[cls.name] = cls  # pylint: disable=protected-access
+    return cls
+
 
 class ComputationSapce(ABC):
     """The abstract base class for compotational spaces.
@@ -16,6 +36,36 @@ class ComputationSapce(ABC):
     regardless of the global setting.
     """
 
+    _registry: Dict[str, Type["ComputationSapce"]] = {}
+
+    @staticmethod
+    def get_comp_space_by_name(name: str) -> Type["ComputationSapce"]:
+        """Get a ComputationSapce by its registered name.
+
+        Args:
+            name (str): The name to probe.
+
+        Returns:
+            Type[ComputationSapce]: The retrieved concrete ComputationSapce.
+        """
+        return ComputationSapce._registry[name]
+
+    def __new__(cls) -> Self:  # TODO: NoReturn should be used, but mypy is not happy.
+        """Raise an error when this class is instantiated.
+
+        Raises:
+            TypeError: This is not the type to instantiate.
+
+        Returns:
+            Self: The method never returns.
+        """
+        raise TypeError("This class cannot be instantiated.")
+
+    # NOTE: Subclasses should implement all the following and should be @final.
+
+    name: str
+    """The name to be registered."""
+
     # TODO: currently only from_log is used. if needed, we can also have to_log, to_lin, from_lin.
     @classmethod
     @abstractmethod