pic

pic

FourierLayer

Bases: Module

Source code in cirkit/backend/torch/parameters/pic.py

class FourierLayer(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        sigma: float | None = 1.0,
        learnable: bool | None = False,
    ):
        super().__init__()
        assert out_features % 2 == 0, "Number of output features must be even."
        self.in_features = in_features
        self.out_features = out_features
        self.sigma = sigma

        coeff = torch.normal(0.0, sigma, (in_features, out_features // 2))
        if learnable:
            self.coeff = nn.Parameter(coeff)
        else:
            self.register_buffer("coeff", coeff)

    def forward(self, z: torch.Tensor):
        z_proj = 2 * torch.pi * z @ self.coeff
        return torch.cat([z_proj.cos(), z_proj.sin()], dim=-1).transpose(-2, -1)

    def extra_repr(self) -> str:
        return f"{self.in_features}, {self.out_features}, sigma={self.sigma}"

coeff = nn.Parameter(coeff) instance-attribute

in_features = in_features instance-attribute

out_features = out_features instance-attribute

sigma = sigma instance-attribute

__init__(in_features, out_features, sigma=1.0, learnable=False)

Source code in cirkit/backend/torch/parameters/pic.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    sigma: float | None = 1.0,
    learnable: bool | None = False,
):
    super().__init__()
    assert out_features % 2 == 0, "Number of output features must be even."
    self.in_features = in_features
    self.out_features = out_features
    self.sigma = sigma

    coeff = torch.normal(0.0, sigma, (in_features, out_features // 2))
    if learnable:
        self.coeff = nn.Parameter(coeff)
    else:
        self.register_buffer("coeff", coeff)

extra_repr()

Source code in cirkit/backend/torch/parameters/pic.py

def extra_repr(self) -> str:
    return f"{self.in_features}, {self.out_features}, sigma={self.sigma}"

forward(z)

Source code in cirkit/backend/torch/parameters/pic.py

def forward(self, z: torch.Tensor):
    z_proj = 2 * torch.pi * z @ self.coeff
    return torch.cat([z_proj.cos(), z_proj.sin()], dim=-1).transpose(-2, -1)

PICInnerNet

Bases: Module

Source code in cirkit/backend/torch/parameters/pic.py

class PICInnerNet(nn.Module):
    def __init__(
        self,
        num_dim: int,
        num_funcs: int,
        perm_dim: tuple[int] | None = None,
        norm_dim: tuple[int] | None = None,
        net_dim: int | None = 64,
        bias: bool | None = False,
        sharing: str | None = "none",
        ff_dim: int | None = None,
        ff_sigma: float | None = 1.0,
        learn_ff: bool | None = False,
        z_quad: torch.Tensor | None = None,
        w_quad: torch.Tensor | None = None,
        tensor_parameter: TorchTensorParameter | None = None,
    ):
        super().__init__()
        assert sharing in ["none", "f", "c"]
        self.num_dim = num_dim
        self.num_funcs = num_funcs
        self.sharing = sharing
        self.perm_dim = (0,) + (tuple(range(1, num_dim + 1)) if perm_dim is None else perm_dim)
        assert self.perm_dim[0] == 0 and set(self.perm_dim) == set(range(num_dim + 1))
        self.norm_dim = (tuple(range(1, num_dim + 1))) if norm_dim is None else norm_dim
        assert 0 not in self.norm_dim and set(self.norm_dim).issubset(self.perm_dim)
        self.eps = np.sqrt(torch.finfo(torch.get_default_dtype()).tiny)
        self.tensor_parameter = tensor_parameter

        assert (z_quad is None) == (w_quad is None), "must both be given or both be None"
        if z_quad is not None:
            self.register_buffer("z_quad", z_quad)
            self.register_buffer("w_quad", w_quad)

        ff_dim = net_dim if ff_dim is None else ff_dim
        inner_conv_groups = 1 if sharing in ["c", "f"] else num_funcs
        last_conv_groups = 1 if sharing == "f" else num_funcs
        self.net = nn.Sequential(
            FourierLayer(num_dim, ff_dim, sigma=ff_sigma, learnable=learn_ff),
            nn.Conv1d(
                inner_conv_groups * ff_dim,
                inner_conv_groups * net_dim,
                1,
                groups=inner_conv_groups,
                bias=bias,
            ),
            nn.Tanh(),
            nn.Conv1d(
                inner_conv_groups * net_dim,
                inner_conv_groups * net_dim,
                1,
                groups=inner_conv_groups,
                bias=bias,
            ),
            nn.Tanh(),
            nn.Conv1d(
                last_conv_groups * net_dim, last_conv_groups, 1, groups=last_conv_groups, bias=bias
            ),
            nn.Softplus(beta=1.0),
        )

        # initialize all heads to be equal when using composite sharing
        if sharing == "c":
            self.net[-2].weight.data = self.net[-2].weight.data[:1].repeat(num_funcs, 1, 1)
            if self.net[-2].bias is not None:
                self.net[-2].bias.data = self.net[-2].bias.data[:1].repeat(num_funcs)

        if tensor_parameter is not None and z_quad is not None:
            with torch.no_grad():
                _ = self()  # initialize tensor_parameter as result of self.forward()

    def forward(
        self,
        z_quad: torch.Tensor | None = None,
        w_quad: torch.Tensor | None = None,
        n_chunks: int | None = 1,
    ):
        z_quad = self.z_quad if z_quad is None else z_quad
        w_quad = self.w_quad if w_quad is None else w_quad
        assert z_quad.ndim == w_quad.ndim == 1 and len(z_quad) == len(w_quad)
        nip = z_quad.numel()  # number of integration points
        self.net[1].groups = 1
        self.net[-2].groups = 1 if self.sharing in ["c", "f"] else self.num_funcs
        z_meshgrid = (
            torch.stack(torch.meshgrid([z_quad] * self.num_dim, indexing="ij")).flatten(1).t()
        )
        logits = (
            torch.cat([self.net(chunk) for chunk in z_meshgrid.chunk(n_chunks, dim=0)], dim=1)
            + self.eps
        )
        # the expand actually does something when self.sharing is 'f'
        logits = (
            logits.expand(self.num_funcs, -1).view(-1, *[nip] * self.num_dim).permute(self.perm_dim)
        )
        w_shape = [nip if i in self.norm_dim else 1 for i in range(self.num_dim + 1)]
        w_meshgrid = (
            torch.stack(torch.meshgrid([w_quad] * len(self.norm_dim), indexing="ij"))
            .prod(0)
            .view(w_shape)
        )
        param = (logits / (logits * w_meshgrid).sum(self.norm_dim, True)) * w_meshgrid
        if self.tensor_parameter is not None:
            param = param.view_as(self.tensor_parameter._ptensor)
            self.tensor_parameter._ptensor = param
        return param

    def __repr__(self):
        return "\n".join(
            [line for line in super().__repr__().split("\n") if "tensor_parameter" not in line]
        )

eps = np.sqrt(torch.finfo(torch.get_default_dtype()).tiny) instance-attribute

net = nn.Sequential(FourierLayer(num_dim, ff_dim, sigma=ff_sigma, learnable=learn_ff), nn.Conv1d(inner_conv_groups * ff_dim, inner_conv_groups * net_dim, 1, groups=inner_conv_groups, bias=bias), nn.Tanh(), nn.Conv1d(inner_conv_groups * net_dim, inner_conv_groups * net_dim, 1, groups=inner_conv_groups, bias=bias), nn.Tanh(), nn.Conv1d(last_conv_groups * net_dim, last_conv_groups, 1, groups=last_conv_groups, bias=bias), nn.Softplus(beta=1.0)) instance-attribute

norm_dim = tuple(range(1, num_dim + 1)) if norm_dim is None else norm_dim instance-attribute

num_dim = num_dim instance-attribute

num_funcs = num_funcs instance-attribute

perm_dim = (0,) + tuple(range(1, num_dim + 1)) if perm_dim is None else perm_dim instance-attribute

sharing = sharing instance-attribute

tensor_parameter = tensor_parameter instance-attribute

__init__(num_dim, num_funcs, perm_dim=None, norm_dim=None, net_dim=64, bias=False, sharing='none', ff_dim=None, ff_sigma=1.0, learn_ff=False, z_quad=None, w_quad=None, tensor_parameter=None)

Source code in cirkit/backend/torch/parameters/pic.py

def __init__(
    self,
    num_dim: int,
    num_funcs: int,
    perm_dim: tuple[int] | None = None,
    norm_dim: tuple[int] | None = None,
    net_dim: int | None = 64,
    bias: bool | None = False,
    sharing: str | None = "none",
    ff_dim: int | None = None,
    ff_sigma: float | None = 1.0,
    learn_ff: bool | None = False,
    z_quad: torch.Tensor | None = None,
    w_quad: torch.Tensor | None = None,
    tensor_parameter: TorchTensorParameter | None = None,
):
    super().__init__()
    assert sharing in ["none", "f", "c"]
    self.num_dim = num_dim
    self.num_funcs = num_funcs
    self.sharing = sharing
    self.perm_dim = (0,) + (tuple(range(1, num_dim + 1)) if perm_dim is None else perm_dim)
    assert self.perm_dim[0] == 0 and set(self.perm_dim) == set(range(num_dim + 1))
    self.norm_dim = (tuple(range(1, num_dim + 1))) if norm_dim is None else norm_dim
    assert 0 not in self.norm_dim and set(self.norm_dim).issubset(self.perm_dim)
    self.eps = np.sqrt(torch.finfo(torch.get_default_dtype()).tiny)
    self.tensor_parameter = tensor_parameter

    assert (z_quad is None) == (w_quad is None), "must both be given or both be None"
    if z_quad is not None:
        self.register_buffer("z_quad", z_quad)
        self.register_buffer("w_quad", w_quad)

    ff_dim = net_dim if ff_dim is None else ff_dim
    inner_conv_groups = 1 if sharing in ["c", "f"] else num_funcs
    last_conv_groups = 1 if sharing == "f" else num_funcs
    self.net = nn.Sequential(
        FourierLayer(num_dim, ff_dim, sigma=ff_sigma, learnable=learn_ff),
        nn.Conv1d(
            inner_conv_groups * ff_dim,
            inner_conv_groups * net_dim,
            1,
            groups=inner_conv_groups,
            bias=bias,
        ),
        nn.Tanh(),
        nn.Conv1d(
            inner_conv_groups * net_dim,
            inner_conv_groups * net_dim,
            1,
            groups=inner_conv_groups,
            bias=bias,
        ),
        nn.Tanh(),
        nn.Conv1d(
            last_conv_groups * net_dim, last_conv_groups, 1, groups=last_conv_groups, bias=bias
        ),
        nn.Softplus(beta=1.0),
    )

    # initialize all heads to be equal when using composite sharing
    if sharing == "c":
        self.net[-2].weight.data = self.net[-2].weight.data[:1].repeat(num_funcs, 1, 1)
        if self.net[-2].bias is not None:
            self.net[-2].bias.data = self.net[-2].bias.data[:1].repeat(num_funcs)

    if tensor_parameter is not None and z_quad is not None:
        with torch.no_grad():
            _ = self()  # initialize tensor_parameter as result of self.forward()

__repr__()

Source code in cirkit/backend/torch/parameters/pic.py

def __repr__(self):
    return "\n".join(
        [line for line in super().__repr__().split("\n") if "tensor_parameter" not in line]
    )

forward(z_quad=None, w_quad=None, n_chunks=1)

Source code in cirkit/backend/torch/parameters/pic.py

def forward(
    self,
    z_quad: torch.Tensor | None = None,
    w_quad: torch.Tensor | None = None,
    n_chunks: int | None = 1,
):
    z_quad = self.z_quad if z_quad is None else z_quad
    w_quad = self.w_quad if w_quad is None else w_quad
    assert z_quad.ndim == w_quad.ndim == 1 and len(z_quad) == len(w_quad)
    nip = z_quad.numel()  # number of integration points
    self.net[1].groups = 1
    self.net[-2].groups = 1 if self.sharing in ["c", "f"] else self.num_funcs
    z_meshgrid = (
        torch.stack(torch.meshgrid([z_quad] * self.num_dim, indexing="ij")).flatten(1).t()
    )
    logits = (
        torch.cat([self.net(chunk) for chunk in z_meshgrid.chunk(n_chunks, dim=0)], dim=1)
        + self.eps
    )
    # the expand actually does something when self.sharing is 'f'
    logits = (
        logits.expand(self.num_funcs, -1).view(-1, *[nip] * self.num_dim).permute(self.perm_dim)
    )
    w_shape = [nip if i in self.norm_dim else 1 for i in range(self.num_dim + 1)]
    w_meshgrid = (
        torch.stack(torch.meshgrid([w_quad] * len(self.norm_dim), indexing="ij"))
        .prod(0)
        .view(w_shape)
    )
    param = (logits / (logits * w_meshgrid).sum(self.norm_dim, True)) * w_meshgrid
    if self.tensor_parameter is not None:
        param = param.view_as(self.tensor_parameter._ptensor)
        self.tensor_parameter._ptensor = param
    return param

PICInputNet

Bases: Module

Source code in cirkit/backend/torch/parameters/pic.py

class PICInputNet(nn.Module):
    def __init__(
        self,
        num_variables: int,
        num_param: int,
        num_channels: bool | None = 1,
        net_dim: int | None = 64,
        bias: bool | None = False,
        sharing: str | None = "none",
        ff_dim: int | None = None,
        ff_sigma: float | None = 1.0,
        learn_ff: bool | None = False,
        z_quad: torch.Tensor | None = None,
        tensor_parameter: TorchTensorParameter | None = None,
        reparam: TorchParameterOp | None = None,
    ):
        super().__init__()
        assert sharing in ["none", "f", "c"]
        self.num_variables = num_variables
        self.num_param = num_param
        self.num_channels = num_channels
        self.sharing = sharing
        self.tensor_parameter = tensor_parameter
        self.reparam = reparam
        if z_quad is not None:
            self.register_buffer("z_quad", z_quad)

        ff_dim = net_dim if ff_dim is None else ff_dim
        inner_conv_groups = num_channels * (1 if sharing in ["f", "c"] else num_variables)
        last_conv_groups = num_channels * (1 if sharing == "f" else num_variables)
        self.net = nn.Sequential(
            FourierLayer(1, ff_dim, sigma=ff_sigma, learnable=learn_ff),
            nn.Conv1d(
                ff_dim * inner_conv_groups,
                net_dim * inner_conv_groups,
                1,
                groups=inner_conv_groups,
                bias=bias,
            ),
            nn.Tanh(),
            nn.Conv1d(
                net_dim * last_conv_groups,
                num_param * last_conv_groups,
                1,
                groups=last_conv_groups,
                bias=bias,
            ),
        )

        # initialize all heads to be equal when using composite sharing
        if sharing == "c":
            self.net[-1].weight.data = (
                self.net[-1].weight.data[: num_param * num_channels].repeat(num_variables, 1, 1)
            )
            if self.net[-1].bias is not None:
                self.net[-1].bias.data = (
                    self.net[-1].bias.data[: num_param * num_channels].repeat(num_variables)
                )

        if tensor_parameter is not None and z_quad is not None:
            with torch.no_grad():
                _ = self()  # initialize tensor_parameter as result of self.forward()

    def forward(self, z_quad: torch.Tensor | None = None, n_chunks: int | None = 1):
        z_quad = self.z_quad if z_quad is None else z_quad
        assert z_quad.ndim == 1
        self.net[1].groups = 1
        self.net[-1].groups = self.num_channels * (
            1 if self.sharing in ["f", "c"] else self.num_variables
        )
        param = torch.cat(
            [self.net(chunk.unsqueeze(1)) for chunk in z_quad.chunk(n_chunks, dim=0)], dim=1
        )
        if self.sharing == "f":
            param = param.unsqueeze(0).expand(self.num_variables, -1, -1)
        param = param.view(
            self.num_variables, self.num_param * self.num_channels, len(z_quad)
        ).transpose(1, 2)
        if self.tensor_parameter is not None:
            param = param.view_as(self.tensor_parameter._ptensor)
            self.tensor_parameter._ptensor = param
        if self.reparam is not None:
            param = self.reparam(param)
        return param

    def __repr__(self):
        return "\n".join(
            [line for line in super().__repr__().split("\n") if "tensor_parameter" not in line]
        )

net = nn.Sequential(FourierLayer(1, ff_dim, sigma=ff_sigma, learnable=learn_ff), nn.Conv1d(ff_dim * inner_conv_groups, net_dim * inner_conv_groups, 1, groups=inner_conv_groups, bias=bias), nn.Tanh(), nn.Conv1d(net_dim * last_conv_groups, num_param * last_conv_groups, 1, groups=last_conv_groups, bias=bias)) instance-attribute

num_channels = num_channels instance-attribute

num_param = num_param instance-attribute

num_variables = num_variables instance-attribute

reparam = reparam instance-attribute

sharing = sharing instance-attribute

tensor_parameter = tensor_parameter instance-attribute

__init__(num_variables, num_param, num_channels=1, net_dim=64, bias=False, sharing='none', ff_dim=None, ff_sigma=1.0, learn_ff=False, z_quad=None, tensor_parameter=None, reparam=None)

Source code in cirkit/backend/torch/parameters/pic.py

def __init__(
    self,
    num_variables: int,
    num_param: int,
    num_channels: bool | None = 1,
    net_dim: int | None = 64,
    bias: bool | None = False,
    sharing: str | None = "none",
    ff_dim: int | None = None,
    ff_sigma: float | None = 1.0,
    learn_ff: bool | None = False,
    z_quad: torch.Tensor | None = None,
    tensor_parameter: TorchTensorParameter | None = None,
    reparam: TorchParameterOp | None = None,
):
    super().__init__()
    assert sharing in ["none", "f", "c"]
    self.num_variables = num_variables
    self.num_param = num_param
    self.num_channels = num_channels
    self.sharing = sharing
    self.tensor_parameter = tensor_parameter
    self.reparam = reparam
    if z_quad is not None:
        self.register_buffer("z_quad", z_quad)

    ff_dim = net_dim if ff_dim is None else ff_dim
    inner_conv_groups = num_channels * (1 if sharing in ["f", "c"] else num_variables)
    last_conv_groups = num_channels * (1 if sharing == "f" else num_variables)
    self.net = nn.Sequential(
        FourierLayer(1, ff_dim, sigma=ff_sigma, learnable=learn_ff),
        nn.Conv1d(
            ff_dim * inner_conv_groups,
            net_dim * inner_conv_groups,
            1,
            groups=inner_conv_groups,
            bias=bias,
        ),
        nn.Tanh(),
        nn.Conv1d(
            net_dim * last_conv_groups,
            num_param * last_conv_groups,
            1,
            groups=last_conv_groups,
            bias=bias,
        ),
    )

    # initialize all heads to be equal when using composite sharing
    if sharing == "c":
        self.net[-1].weight.data = (
            self.net[-1].weight.data[: num_param * num_channels].repeat(num_variables, 1, 1)
        )
        if self.net[-1].bias is not None:
            self.net[-1].bias.data = (
                self.net[-1].bias.data[: num_param * num_channels].repeat(num_variables)
            )

    if tensor_parameter is not None and z_quad is not None:
        with torch.no_grad():
            _ = self()  # initialize tensor_parameter as result of self.forward()

__repr__()

Source code in cirkit/backend/torch/parameters/pic.py

def __repr__(self):
    return "\n".join(
        [line for line in super().__repr__().split("\n") if "tensor_parameter" not in line]
    )

forward(z_quad=None, n_chunks=1)

Source code in cirkit/backend/torch/parameters/pic.py

def forward(self, z_quad: torch.Tensor | None = None, n_chunks: int | None = 1):
    z_quad = self.z_quad if z_quad is None else z_quad
    assert z_quad.ndim == 1
    self.net[1].groups = 1
    self.net[-1].groups = self.num_channels * (
        1 if self.sharing in ["f", "c"] else self.num_variables
    )
    param = torch.cat(
        [self.net(chunk.unsqueeze(1)) for chunk in z_quad.chunk(n_chunks, dim=0)], dim=1
    )
    if self.sharing == "f":
        param = param.unsqueeze(0).expand(self.num_variables, -1, -1)
    param = param.view(
        self.num_variables, self.num_param * self.num_channels, len(z_quad)
    ).transpose(1, 2)
    if self.tensor_parameter is not None:
        param = param.view_as(self.tensor_parameter._ptensor)
        self.tensor_parameter._ptensor = param
    if self.reparam is not None:
        param = self.reparam(param)
    return param

TorchCPTLayer

Bases: TorchInnerLayer

The Candecomp transposed (CP-T) layer, which is the fusion of a sum layer and a Hadamard layer.

Source code in cirkit/backend/torch/layers/optimized.py

class TorchCPTLayer(TorchInnerLayer):
    """The Candecomp transposed (CP-T) layer, which is the fusion of a sum layer and a Hadamard
    layer.
    """

    def __init__(
        self,
        num_input_units: int,
        num_output_units: int,
        arity: int = 2,
        *,
        weight: TorchParameter,
        semiring: Semiring | None = None,
        num_folds: int = 1,
    ):
        """Initialize a CP-T layer.

        Args:
            num_input_units: The number of input units.
            num_output_units: The number of output units.
            arity: The arity of the layer, must be 2. Defaults to 2.
            weight: The weight parameter, which must have shape $(F, K_o, K_i)$,
                where $F$ is the number of folds, $K_o$ is the number output units,
                and $K_i$ is the number of input units.

        Raises:
            ValueError: If the number of input and output units are incompatible with the
                shape of the weight parameter.
        """
        super().__init__(
            num_input_units,
            num_output_units,
            arity=arity,
            semiring=semiring,
            num_folds=num_folds,
        )
        if not self._valid_weight_shape(weight):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._weight_shape} for 'weight', found"
                f"{weight.num_folds} and {weight.shape}, respectively"
            )
        self.weight = weight

    def _valid_weight_shape(self, w: TorchParameter) -> bool:
        if w.num_folds != self.num_folds:
            return False
        return w.shape == self._weight_shape

    @property
    def _weight_shape(self) -> tuple[int, ...]:
        return self.num_output_units, self.num_input_units

    @property
    def config(self) -> Mapping[str, Any]:
        return {
            "num_input_units": self.num_input_units,
            "num_output_units": self.num_output_units,
            "arity": self.arity,
        }

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        return {"weight": self.weight}

    def forward(self, x: Tensor) -> Tensor:
        # x: (F, B, Ki)
        x = self.semiring.prod(x, dim=1, keepdim=False)
        # weight: (F, Ko, Ki)
        weight = self.weight()
        return self.semiring.einsum(
            "fbi,foi->fbo", inputs=(x,), operands=(weight,), dim=-1, keepdim=True
        )

    def sample(self, x: Tensor) -> tuple[Tensor, Tensor]:
        weight = self.weight()
        negative = torch.any(weight < 0.0)
        if negative:
            raise ValueError("Sampling only works with positive weights")
        normalized = torch.allclose(torch.sum(weight, dim=-1), torch.ones(1, device=weight.device))
        if not normalized:
            raise ValueError("Sampling only works with a normalized parametrization")

        # x: (F, H, C, K, num_samples, D)
        x = torch.sum(x, dim=1, keepdim=True)  # (F, H=1, C, K, num_samples, D)

        c = x.shape[2]
        d = x.shape[-1]
        num_samples = x.shape[-2]

        # mixing_distribution: (F, O, K)
        mixing_distribution = torch.distributions.Categorical(probs=weight)

        mixing_samples = mixing_distribution.sample((num_samples,))
        mixing_samples = E.rearrange(mixing_samples, "n f o -> f o n")
        mixing_indices = E.repeat(mixing_samples, "f o n -> f a c o n d", a=1, c=c, d=d)

        x = torch.gather(x, dim=-3, index=mixing_indices)
        x = x[:, 0]
        return x, mixing_samples

_weight_shape property

config property

params property

weight = weight instance-attribute

__init__(num_input_units, num_output_units, arity=2, *, weight, semiring=None, num_folds=1)

Initialize a CP-T layer.

Parameters:

Name	Type	Description	Default
num_input_units	int	The number of input units.	required
num_output_units	int	The number of output units.	required
arity	int	The arity of the layer, must be 2. Defaults to 2.	2
weight	TorchParameter	The weight parameter, which must have shape \((F, K_o, K_i)\), where \(F\) is the number of folds, \(K_o\) is the number output units, and \(K_i\) is the number of input units.	required

Raises:

Type	Description
ValueError	If the number of input and output units are incompatible with the shape of the weight parameter.

Source code in cirkit/backend/torch/layers/optimized.py

def __init__(
    self,
    num_input_units: int,
    num_output_units: int,
    arity: int = 2,
    *,
    weight: TorchParameter,
    semiring: Semiring | None = None,
    num_folds: int = 1,
):
    """Initialize a CP-T layer.

    Args:
        num_input_units: The number of input units.
        num_output_units: The number of output units.
        arity: The arity of the layer, must be 2. Defaults to 2.
        weight: The weight parameter, which must have shape $(F, K_o, K_i)$,
            where $F$ is the number of folds, $K_o$ is the number output units,
            and $K_i$ is the number of input units.

    Raises:
        ValueError: If the number of input and output units are incompatible with the
            shape of the weight parameter.
    """
    super().__init__(
        num_input_units,
        num_output_units,
        arity=arity,
        semiring=semiring,
        num_folds=num_folds,
    )
    if not self._valid_weight_shape(weight):
        raise ValueError(
            f"Expected number of folds {self.num_folds} "
            f"and shape {self._weight_shape} for 'weight', found"
            f"{weight.num_folds} and {weight.shape}, respectively"
        )
    self.weight = weight

_valid_weight_shape(w)

Source code in cirkit/backend/torch/layers/optimized.py

def _valid_weight_shape(self, w: TorchParameter) -> bool:
    if w.num_folds != self.num_folds:
        return False
    return w.shape == self._weight_shape

forward(x)

Source code in cirkit/backend/torch/layers/optimized.py

def forward(self, x: Tensor) -> Tensor:
    # x: (F, B, Ki)
    x = self.semiring.prod(x, dim=1, keepdim=False)
    # weight: (F, Ko, Ki)
    weight = self.weight()
    return self.semiring.einsum(
        "fbi,foi->fbo", inputs=(x,), operands=(weight,), dim=-1, keepdim=True
    )

sample(x)

Source code in cirkit/backend/torch/layers/optimized.py

def sample(self, x: Tensor) -> tuple[Tensor, Tensor]:
    weight = self.weight()
    negative = torch.any(weight < 0.0)
    if negative:
        raise ValueError("Sampling only works with positive weights")
    normalized = torch.allclose(torch.sum(weight, dim=-1), torch.ones(1, device=weight.device))
    if not normalized:
        raise ValueError("Sampling only works with a normalized parametrization")

    # x: (F, H, C, K, num_samples, D)
    x = torch.sum(x, dim=1, keepdim=True)  # (F, H=1, C, K, num_samples, D)

    c = x.shape[2]
    d = x.shape[-1]
    num_samples = x.shape[-2]

    # mixing_distribution: (F, O, K)
    mixing_distribution = torch.distributions.Categorical(probs=weight)

    mixing_samples = mixing_distribution.sample((num_samples,))
    mixing_samples = E.rearrange(mixing_samples, "n f o -> f o n")
    mixing_indices = E.repeat(mixing_samples, "f o n -> f a c o n d", a=1, c=c, d=d)

    x = torch.gather(x, dim=-3, index=mixing_indices)
    x = x[:, 0]
    return x, mixing_samples

TorchCategoricalLayer

Bases: TorchExpFamilyLayer

The Categorical distribution layer, parameterized by either probabilities or logits.

Source code in cirkit/backend/torch/layers/input.py

class TorchCategoricalLayer(TorchExpFamilyLayer):
    """The Categorical distribution layer, parameterized by either probabilities or logits."""

    # pylint: disable-next=too-many-arguments
    def __init__(
        self,
        scope_idx: Tensor,
        num_output_units: int,
        num_channels: int = 1,
        *,
        num_categories: int = 2,
        probs: TorchParameter | None = None,
        logits: TorchParameter | None = None,
        semiring: Semiring | None = None,
    ) -> None:
        """Initialize a Categorical layer.

        Args:
            scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
                $D$ is the number of variables on which the input layers in each fold are defined on.
                Alternatively, a tensor of shape $(D,)$ can be specified, which will be interpreted
                as a tensor of shape $(1, D)$, i.e., with $F = 1$.
            num_output_units: The number of output units.
            num_channels: The number of channels.
            num_categories: The number of categories for Categorical distribution.
            probs: The probabilities parameter of shape $(F, K, C, V)$, where $K$ is the number of
                output units, $C$ is the number of channels, and $V$ is the number of categories.
            logits: The logits parameter of shape $(F, K, C, V)$, where $K$ is the number of
                output units, $C$ is the number of channels, and $V$ is the number of categories.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

        Raises:
            ValueError: If the scope contains more than one variable.
            ValueError: If the number of categories is negative.
            ValueError: If both the probs and logits parameters are provided, or none of them.
            ValueError: If the parameter's shape is incorrect.
        """
        num_variables = scope_idx.shape[-1]
        if num_variables != 1:
            raise ValueError("The Gaussian layer encodes a univariate distribution")
        if num_categories <= 0:
            raise ValueError(
                "The number of categories for Categorical distribution must be positive"
            )
        super().__init__(
            scope_idx,
            num_output_units,
            num_channels=num_channels,
            semiring=semiring,
        )
        self.num_categories = num_categories
        if not ((logits is None) ^ (probs is None)):
            raise ValueError("Exactly one between 'logits' and 'probs' must be specified")
        if logits is None:
            assert probs is not None
            if not self._valid_parameter_shape(probs):
                raise ValueError(
                    f"Expected number of folds {self.num_folds} "
                    f"and shape {self._probs_logits_shape} for 'probs', found"
                    f"{probs.num_folds} and {probs.shape}, respectively"
                )
        elif not self._valid_parameter_shape(logits):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._probs_logits_shape} for 'logits', found"
                f"{logits.num_folds} and {logits.shape}, respectively"
            )
        self.probs = probs
        self.logits = logits

    def _valid_parameter_shape(self, p: TorchParameter) -> bool:
        if p.num_folds != self.num_folds:
            return False
        return p.shape == self._probs_logits_shape

    @property
    def _probs_logits_shape(self) -> tuple[int, ...]:
        return self.num_output_units, self.num_channels, self.num_categories

    @property
    def config(self) -> Mapping[str, Any]:
        return {
            "num_output_units": self.num_output_units,
            "num_channels": self.num_channels,
            "num_categories": self.num_categories,
        }

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        if self.logits is None:
            return {"probs": self.probs}
        return {"logits": self.logits}

    def log_unnormalized_likelihood(self, x: Tensor) -> Tensor:
        if x.is_floating_point():
            x = x.long()  # The input to Categorical should be discrete
        # x: (F, C, B, 1) -> (F, C, B)
        x = x.squeeze(dim=3)
        # logits: (F, K, C, N)
        logits = torch.log(self.probs()) if self.logits is None else self.logits()
        if self.num_channels == 1:
            idx_fold = torch.arange(self.num_folds)
            x = logits[:, :, 0][idx_fold[:, None], :, x[:, 0]]
        else:
            idx_fold = torch.arange(self.num_folds)[:, None, None]
            idx_channel = torch.arange(self.num_channels)[None, :, None]
            x = torch.sum(logits[idx_fold, :, idx_channel, x], dim=1)
        return x

    def log_partition_function(self) -> Tensor:
        if self.logits is None:
            return torch.zeros(
                size=(self.num_folds, 1, self.num_output_units), device=self.probs.device
            )
        logits = self.logits()
        return torch.sum(torch.logsumexp(logits, dim=3), dim=2).unsqueeze(dim=1)

    def sample(self, num_samples: int = 1) -> Tensor:
        logits = torch.log(self.probs()) if self.logits is None else self.logits()
        dist = distributions.Categorical(logits=logits)
        samples = dist.sample((num_samples,))  # (N, F, K, C)
        samples = samples.permute(1, 3, 2, 0)  # (F, C, K, N)
        return samples

_probs_logits_shape property

config property

logits = logits instance-attribute

num_categories = num_categories instance-attribute

params property

probs = probs instance-attribute

__init__(scope_idx, num_output_units, num_channels=1, *, num_categories=2, probs=None, logits=None, semiring=None)

Initialize a Categorical layer.

Parameters:

Name	Type	Description	Default
scope_idx	Tensor	A tensor of shape \((F, D)\), where \(F\) is the number of folds, and \(D\) is the number of variables on which the input layers in each fold are defined on. Alternatively, a tensor of shape \((D,)\) can be specified, which will be interpreted as a tensor of shape \((1, D)\), i.e., with \(F = 1\).	required
num_output_units	int	The number of output units.	required
num_channels	int	The number of channels.	1
num_categories	int	The number of categories for Categorical distribution.	2
probs	TorchParameter | None	The probabilities parameter of shape \((F, K, C, V)\), where \(K\) is the number of output units, \(C\) is the number of channels, and \(V\) is the number of categories.	None
logits	TorchParameter | None	The logits parameter of shape \((F, K, C, V)\), where \(K\) is the number of output units, \(C\) is the number of channels, and \(V\) is the number of categories.	None
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None

Raises:

Type	Description
ValueError	If the scope contains more than one variable.
ValueError	If the number of categories is negative.
ValueError	If both the probs and logits parameters are provided, or none of them.
ValueError	If the parameter's shape is incorrect.

Source code in cirkit/backend/torch/layers/input.py

def __init__(
    self,
    scope_idx: Tensor,
    num_output_units: int,
    num_channels: int = 1,
    *,
    num_categories: int = 2,
    probs: TorchParameter | None = None,
    logits: TorchParameter | None = None,
    semiring: Semiring | None = None,
) -> None:
    """Initialize a Categorical layer.

    Args:
        scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
            $D$ is the number of variables on which the input layers in each fold are defined on.
            Alternatively, a tensor of shape $(D,)$ can be specified, which will be interpreted
            as a tensor of shape $(1, D)$, i.e., with $F = 1$.
        num_output_units: The number of output units.
        num_channels: The number of channels.
        num_categories: The number of categories for Categorical distribution.
        probs: The probabilities parameter of shape $(F, K, C, V)$, where $K$ is the number of
            output units, $C$ is the number of channels, and $V$ is the number of categories.
        logits: The logits parameter of shape $(F, K, C, V)$, where $K$ is the number of
            output units, $C$ is the number of channels, and $V$ is the number of categories.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

    Raises:
        ValueError: If the scope contains more than one variable.
        ValueError: If the number of categories is negative.
        ValueError: If both the probs and logits parameters are provided, or none of them.
        ValueError: If the parameter's shape is incorrect.
    """
    num_variables = scope_idx.shape[-1]
    if num_variables != 1:
        raise ValueError("The Gaussian layer encodes a univariate distribution")
    if num_categories <= 0:
        raise ValueError(
            "The number of categories for Categorical distribution must be positive"
        )
    super().__init__(
        scope_idx,
        num_output_units,
        num_channels=num_channels,
        semiring=semiring,
    )
    self.num_categories = num_categories
    if not ((logits is None) ^ (probs is None)):
        raise ValueError("Exactly one between 'logits' and 'probs' must be specified")
    if logits is None:
        assert probs is not None
        if not self._valid_parameter_shape(probs):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._probs_logits_shape} for 'probs', found"
                f"{probs.num_folds} and {probs.shape}, respectively"
            )
    elif not self._valid_parameter_shape(logits):
        raise ValueError(
            f"Expected number of folds {self.num_folds} "
            f"and shape {self._probs_logits_shape} for 'logits', found"
            f"{logits.num_folds} and {logits.shape}, respectively"
        )
    self.probs = probs
    self.logits = logits

_valid_parameter_shape(p)

Source code in cirkit/backend/torch/layers/input.py

def _valid_parameter_shape(self, p: TorchParameter) -> bool:
    if p.num_folds != self.num_folds:
        return False
    return p.shape == self._probs_logits_shape

log_partition_function()

Source code in cirkit/backend/torch/layers/input.py

def log_partition_function(self) -> Tensor:
    if self.logits is None:
        return torch.zeros(
            size=(self.num_folds, 1, self.num_output_units), device=self.probs.device
        )
    logits = self.logits()
    return torch.sum(torch.logsumexp(logits, dim=3), dim=2).unsqueeze(dim=1)

log_unnormalized_likelihood(x)

Source code in cirkit/backend/torch/layers/input.py

def log_unnormalized_likelihood(self, x: Tensor) -> Tensor:
    if x.is_floating_point():
        x = x.long()  # The input to Categorical should be discrete
    # x: (F, C, B, 1) -> (F, C, B)
    x = x.squeeze(dim=3)
    # logits: (F, K, C, N)
    logits = torch.log(self.probs()) if self.logits is None else self.logits()
    if self.num_channels == 1:
        idx_fold = torch.arange(self.num_folds)
        x = logits[:, :, 0][idx_fold[:, None], :, x[:, 0]]
    else:
        idx_fold = torch.arange(self.num_folds)[:, None, None]
        idx_channel = torch.arange(self.num_channels)[None, :, None]
        x = torch.sum(logits[idx_fold, :, idx_channel, x], dim=1)
    return x

sample(num_samples=1)

Source code in cirkit/backend/torch/layers/input.py

def sample(self, num_samples: int = 1) -> Tensor:
    logits = torch.log(self.probs()) if self.logits is None else self.logits()
    dist = distributions.Categorical(logits=logits)
    samples = dist.sample((num_samples,))  # (N, F, K, C)
    samples = samples.permute(1, 3, 2, 0)  # (F, C, K, N)
    return samples

TorchExpFamilyLayer

Bases: TorchInputFunctionLayer, ABC

The abstract base class for exponential family distribution layers. An input layer that is an exponential family distribution must define two methods. The first one is the log_unnormalized_likelihood, used to compute the possibly-unnormalized log-likelihood. The second one is the log_partition_function, used to compute the logarithm of the partition function.

Source code in cirkit/backend/torch/layers/input.py

class TorchExpFamilyLayer(TorchInputFunctionLayer, ABC):
    """The abstract base class for exponential family distribution layers.
    An input layer that is an exponential family distribution must define two methods.
    The first one is the ```log_unnormalized_likelihood```, used to compute the
    possibly-unnormalized log-likelihood. The second one is the ```log_partition_function```,
    used to compute the logarithm of the partition function."""

    def forward(self, x: Tensor) -> Tensor:
        x = self.log_unnormalized_likelihood(x)
        return self.semiring.map_from(x, LSESumSemiring)

    def integrate(self) -> Tensor:
        log_partition = self.log_partition_function()
        return self.semiring.map_from(log_partition, LSESumSemiring)

    @abstractmethod
    def log_unnormalized_likelihood(self, x: Tensor) -> Tensor:
        """Compute the (possibly unnormalized) log-likelihood of the given inputs.

        Args:
            x: The input tensor.

        Returns:
            Tensor: The (possibly unnormalized) log-likelihood as a tensor of shape $(F, K)$,
                where $F$ is the number of folds and $K$ is the number of output units.
        """

    @abstractmethod
    def log_partition_function(self) -> Tensor:
        """Compute the logarithm of the partition function of the layer.

        Returns:
            Tensor: The logarithm of the partition function as a tensor of shape $(F, K)$,
                where $F$ is the number of folds and $K$ is the number of output units.
                Note that it will be a tensor of zeros if the layer encodes already normalized
                exponential family distributions.
        """

forward(x)

Source code in cirkit/backend/torch/layers/input.py

def forward(self, x: Tensor) -> Tensor:
    x = self.log_unnormalized_likelihood(x)
    return self.semiring.map_from(x, LSESumSemiring)

integrate()

Source code in cirkit/backend/torch/layers/input.py

def integrate(self) -> Tensor:
    log_partition = self.log_partition_function()
    return self.semiring.map_from(log_partition, LSESumSemiring)

log_partition_function() abstractmethod

Compute the logarithm of the partition function of the layer.

Returns:

Name	Type	Description
Tensor	Tensor	The logarithm of the partition function as a tensor of shape \((F, K)\), where \(F\) is the number of folds and \(K\) is the number of output units. Note that it will be a tensor of zeros if the layer encodes already normalized exponential family distributions.

Source code in cirkit/backend/torch/layers/input.py

@abstractmethod
def log_partition_function(self) -> Tensor:
    """Compute the logarithm of the partition function of the layer.

    Returns:
        Tensor: The logarithm of the partition function as a tensor of shape $(F, K)$,
            where $F$ is the number of folds and $K$ is the number of output units.
            Note that it will be a tensor of zeros if the layer encodes already normalized
            exponential family distributions.
    """

log_unnormalized_likelihood(x) abstractmethod

Compute the (possibly unnormalized) log-likelihood of the given inputs.

Parameters:

Name	Type	Description	Default
x	Tensor	The input tensor.	required

Returns:

Name	Type	Description
Tensor	Tensor	The (possibly unnormalized) log-likelihood as a tensor of shape \((F, K)\), where \(F\) is the number of folds and \(K\) is the number of output units.

Source code in cirkit/backend/torch/layers/input.py

@abstractmethod
def log_unnormalized_likelihood(self, x: Tensor) -> Tensor:
    """Compute the (possibly unnormalized) log-likelihood of the given inputs.

    Args:
        x: The input tensor.

    Returns:
        Tensor: The (possibly unnormalized) log-likelihood as a tensor of shape $(F, K)$,
            where $F$ is the number of folds and $K$ is the number of output units.
    """

TorchGaussianLayer

Bases: TorchExpFamilyLayer

The Gaussian distribution layer. Optionally, this layer can encode unnormalized Gaussian distributions with the spefication of a log-partition function parameter.

Source code in cirkit/backend/torch/layers/input.py

class TorchGaussianLayer(TorchExpFamilyLayer):
    """The Gaussian distribution layer. Optionally, this layer can encode unnormalized Gaussian
    distributions with the spefication of a log-partition function parameter."""

    def __init__(
        self,
        scope_idx: Tensor,
        num_output_units: int,
        num_channels: int = 1,
        *,
        mean: TorchParameter,
        stddev: TorchParameter,
        log_partition: TorchParameter | None = None,
        semiring: Semiring | None = None,
    ) -> None:
        r"""Initialize a Gaussian layer.

        Args:
            scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
                $D$ is the number of variables on which the input layers in each fold are defined on.
                Alternatively, a tensor of shape $(D,)$ can be specified, which will be interpreted
                as a tensor of shape $(1, D)$, i.e., with $F = 1$.
            num_output_units: The number of output units.
            num_channels: The number of channels.
            mean: The mean parameter, having shape $(F, K, C)$, where $K$ is the number of
                output units and $C$ is the number of channels.
            stddev: The standard deviation parameter, having shape $(F, K, C)$, where $K$ is the
                number of output units and $C$ is the number of channels.
            log_partition: An optional parameter of shape $(F, K, C)$, encoding the log-partition.
                function. If this is not None, then the Gaussian layer encodes unnormalized
                Gaussian likelihoods, which are then normalized with the given log-partition
                function.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

        Raises:
            ValueError: If the scope contains more than one variable.
            ValueError: If the mean and standard deviation parameter shapes are incorrect.
            ValueError: If the log-partition function parameter shape is incorrect.
        """
        num_variables = scope_idx.shape[-1]
        if num_variables != 1:
            raise ValueError("The Gaussian layer encodes a univariate distribution")
        super().__init__(
            scope_idx,
            num_output_units,
            num_channels=num_channels,
            semiring=semiring,
        )
        if not self._valid_mean_stddev_shape(mean):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._mean_stddev_shape} for 'mean', found"
                f"{mean.num_folds} and {mean.shape}, respectively"
            )
        if not self._valid_mean_stddev_shape(stddev):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._mean_stddev_shape} for 'stddev', found"
                f"{stddev.num_folds} and {stddev.shape}, respectively"
            )
        if log_partition is not None and not self._valid_log_partition_shape(log_partition):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._log_partition_shape} for 'log_partition', found"
                f"{log_partition.num_folds} and {log_partition.shape}, respectively"
            )
        self.mean = mean
        self.stddev = stddev
        self.log_partition = log_partition

    def _valid_mean_stddev_shape(self, p: TorchParameter) -> bool:
        if p.num_folds != self.num_folds:
            return False
        return p.shape == self._mean_stddev_shape

    def _valid_log_partition_shape(self, log_partition: TorchParameter) -> bool:
        if log_partition.num_folds != self.num_folds:
            return False
        return log_partition.shape == self._log_partition_shape

    @property
    def _mean_stddev_shape(self) -> tuple[int, ...]:
        return self.num_output_units, self.num_channels

    @property
    def _log_partition_shape(self) -> tuple[int, ...]:
        return self.num_output_units, self.num_channels

    @property
    def config(self) -> Mapping[str, Any]:
        return {"num_output_units": self.num_output_units, "num_channels": self.num_channels}

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        params = {"mean": self.mean, "stddev": self.stddev}
        if self.log_partition is not None:
            params["log_partition"] = self.log_partition
        return params

    def log_unnormalized_likelihood(self, x: Tensor) -> Tensor:
        mean = self.mean().unsqueeze(dim=1)  # (F, 1, K, C)
        stddev = self.stddev().unsqueeze(dim=1)  # (F, 1, K, C)
        x = x.permute(0, 2, 3, 1)  # (F, C, B, 1) -> (F, B, 1, C)
        x = distributions.Normal(loc=mean, scale=stddev).log_prob(x)  # (F, B, K, C)
        x = torch.sum(x, dim=3)  # (F, B, K)
        if self.log_partition is not None:
            log_partition = self.log_partition()  # (F, K, C)
            x = x + torch.sum(log_partition, dim=2).unsqueeze(dim=1)
        return x

    def log_partition_function(self) -> Tensor:
        if self.log_partition is None:
            return torch.zeros(
                size=(self.num_folds, 1, self.num_output_units), device=self.mean.device
            )
        log_partition = self.log_partition()  # (F, K, C)
        return torch.sum(log_partition, dim=2).unsqueeze(dim=1)

    def sample(self, num_samples: int = 1) -> Tensor:
        dist = distributions.Normal(loc=self.mean(), scale=self.stddev())
        samples = dist.sample((num_samples,))  # (N, F, K, C)
        samples = samples.permute(1, 3, 2, 0)  # (F, C, K, N)
        return samples

_log_partition_shape property

_mean_stddev_shape property

config property

log_partition = log_partition instance-attribute

mean = mean instance-attribute

params property

stddev = stddev instance-attribute

__init__(scope_idx, num_output_units, num_channels=1, *, mean, stddev, log_partition=None, semiring=None)

Initialize a Gaussian layer.

Parameters:

Name	Type	Description	Default
scope_idx	Tensor	A tensor of shape \((F, D)\), where \(F\) is the number of folds, and \(D\) is the number of variables on which the input layers in each fold are defined on. Alternatively, a tensor of shape \((D,)\) can be specified, which will be interpreted as a tensor of shape \((1, D)\), i.e., with \(F = 1\).	required
num_output_units	int	The number of output units.	required
num_channels	int	The number of channels.	1
mean	TorchParameter	The mean parameter, having shape \((F, K, C)\), where \(K\) is the number of output units and \(C\) is the number of channels.	required
stddev	TorchParameter	The standard deviation parameter, having shape \((F, K, C)\), where \(K\) is the number of output units and \(C\) is the number of channels.	required
log_partition	TorchParameter | None	An optional parameter of shape \((F, K, C)\), encoding the log-partition. function. If this is not None, then the Gaussian layer encodes unnormalized Gaussian likelihoods, which are then normalized with the given log-partition function.	None
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None

Raises:

Type	Description
ValueError	If the scope contains more than one variable.
ValueError	If the mean and standard deviation parameter shapes are incorrect.
ValueError	If the log-partition function parameter shape is incorrect.

Source code in cirkit/backend/torch/layers/input.py

def __init__(
    self,
    scope_idx: Tensor,
    num_output_units: int,
    num_channels: int = 1,
    *,
    mean: TorchParameter,
    stddev: TorchParameter,
    log_partition: TorchParameter | None = None,
    semiring: Semiring | None = None,
) -> None:
    r"""Initialize a Gaussian layer.

    Args:
        scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
            $D$ is the number of variables on which the input layers in each fold are defined on.
            Alternatively, a tensor of shape $(D,)$ can be specified, which will be interpreted
            as a tensor of shape $(1, D)$, i.e., with $F = 1$.
        num_output_units: The number of output units.
        num_channels: The number of channels.
        mean: The mean parameter, having shape $(F, K, C)$, where $K$ is the number of
            output units and $C$ is the number of channels.
        stddev: The standard deviation parameter, having shape $(F, K, C)$, where $K$ is the
            number of output units and $C$ is the number of channels.
        log_partition: An optional parameter of shape $(F, K, C)$, encoding the log-partition.
            function. If this is not None, then the Gaussian layer encodes unnormalized
            Gaussian likelihoods, which are then normalized with the given log-partition
            function.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

    Raises:
        ValueError: If the scope contains more than one variable.
        ValueError: If the mean and standard deviation parameter shapes are incorrect.
        ValueError: If the log-partition function parameter shape is incorrect.
    """
    num_variables = scope_idx.shape[-1]
    if num_variables != 1:
        raise ValueError("The Gaussian layer encodes a univariate distribution")
    super().__init__(
        scope_idx,
        num_output_units,
        num_channels=num_channels,
        semiring=semiring,
    )
    if not self._valid_mean_stddev_shape(mean):
        raise ValueError(
            f"Expected number of folds {self.num_folds} "
            f"and shape {self._mean_stddev_shape} for 'mean', found"
            f"{mean.num_folds} and {mean.shape}, respectively"
        )
    if not self._valid_mean_stddev_shape(stddev):
        raise ValueError(
            f"Expected number of folds {self.num_folds} "
            f"and shape {self._mean_stddev_shape} for 'stddev', found"
            f"{stddev.num_folds} and {stddev.shape}, respectively"
        )
    if log_partition is not None and not self._valid_log_partition_shape(log_partition):
        raise ValueError(
            f"Expected number of folds {self.num_folds} "
            f"and shape {self._log_partition_shape} for 'log_partition', found"
            f"{log_partition.num_folds} and {log_partition.shape}, respectively"
        )
    self.mean = mean
    self.stddev = stddev
    self.log_partition = log_partition

_valid_log_partition_shape(log_partition)

Source code in cirkit/backend/torch/layers/input.py

def _valid_log_partition_shape(self, log_partition: TorchParameter) -> bool:
    if log_partition.num_folds != self.num_folds:
        return False
    return log_partition.shape == self._log_partition_shape

_valid_mean_stddev_shape(p)

Source code in cirkit/backend/torch/layers/input.py

def _valid_mean_stddev_shape(self, p: TorchParameter) -> bool:
    if p.num_folds != self.num_folds:
        return False
    return p.shape == self._mean_stddev_shape

log_partition_function()

Source code in cirkit/backend/torch/layers/input.py

def log_partition_function(self) -> Tensor:
    if self.log_partition is None:
        return torch.zeros(
            size=(self.num_folds, 1, self.num_output_units), device=self.mean.device
        )
    log_partition = self.log_partition()  # (F, K, C)
    return torch.sum(log_partition, dim=2).unsqueeze(dim=1)

log_unnormalized_likelihood(x)

Source code in cirkit/backend/torch/layers/input.py

def log_unnormalized_likelihood(self, x: Tensor) -> Tensor:
    mean = self.mean().unsqueeze(dim=1)  # (F, 1, K, C)
    stddev = self.stddev().unsqueeze(dim=1)  # (F, 1, K, C)
    x = x.permute(0, 2, 3, 1)  # (F, C, B, 1) -> (F, B, 1, C)
    x = distributions.Normal(loc=mean, scale=stddev).log_prob(x)  # (F, B, K, C)
    x = torch.sum(x, dim=3)  # (F, B, K)
    if self.log_partition is not None:
        log_partition = self.log_partition()  # (F, K, C)
        x = x + torch.sum(log_partition, dim=2).unsqueeze(dim=1)
    return x

sample(num_samples=1)

Source code in cirkit/backend/torch/layers/input.py

def sample(self, num_samples: int = 1) -> Tensor:
    dist = distributions.Normal(loc=self.mean(), scale=self.stddev())
    samples = dist.sample((num_samples,))  # (N, F, K, C)
    samples = samples.permute(1, 3, 2, 0)  # (F, C, K, N)
    return samples

TorchHadamardLayer

Bases: TorchInnerLayer

The Hadamard product layer, which computes an element-wise (or Hadamard) product of the input vectors it receives as inputs. See the symbolic HadamardLayer for more details.

Source code in cirkit/backend/torch/layers/inner.py

class TorchHadamardLayer(TorchInnerLayer):
    """The Hadamard product layer, which computes an element-wise (or Hadamard) product of
    the input vectors it receives as inputs.
    See the symbolic [HadamardLayer][cirkit.symbolic.layers.HadamardLayer] for more details.
    """

    def __init__(
        self,
        num_input_units: int,
        arity: int = 2,
        *,
        semiring: Semiring | None = None,
        num_folds: int = 1,
    ):
        """Initialize a Hadamard product layer.

        Args:
            num_input_units: The number of input units, which is equal to the number of
                output units.
            arity: The arity of the layer.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
            num_folds: The number of channels.

        Raises:
            ValueError: If the arity is not at least 2.
            ValueError: If the number of input units is not the same as the number of output units.
        """
        if arity < 2:
            raise ValueError("The arity should be at least 2")
        super().__init__(
            num_input_units, num_input_units, arity=arity, semiring=semiring, num_folds=num_folds
        )

    @property
    def config(self) -> Mapping[str, Any]:
        return {
            "num_input_units": self.num_input_units,
            "arity": self.arity,
        }

    def forward(self, x: Tensor) -> Tensor:
        return self.semiring.prod(x, dim=1, keepdim=False)  # shape (F, H, B, K) -> (F, B, K).

    def sample(self, x: Tensor) -> tuple[Tensor, None]:
        # Concatenate samples over disjoint variables through a sum
        # x: (F, H, C, K, num_samples, D)
        x = torch.sum(x, dim=1)  # (F, C, K, num_samples, D)
        return x, None

config property

__init__(num_input_units, arity=2, *, semiring=None, num_folds=1)

Initialize a Hadamard product layer.

Parameters:

Name	Type	Description	Default
num_input_units	int	The number of input units, which is equal to the number of output units.	required
arity	int	The arity of the layer.	2
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None
num_folds	int	The number of channels.	1

Raises:

Type	Description
ValueError	If the arity is not at least 2.
ValueError	If the number of input units is not the same as the number of output units.

Source code in cirkit/backend/torch/layers/inner.py

def __init__(
    self,
    num_input_units: int,
    arity: int = 2,
    *,
    semiring: Semiring | None = None,
    num_folds: int = 1,
):
    """Initialize a Hadamard product layer.

    Args:
        num_input_units: The number of input units, which is equal to the number of
            output units.
        arity: The arity of the layer.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
        num_folds: The number of channels.

    Raises:
        ValueError: If the arity is not at least 2.
        ValueError: If the number of input units is not the same as the number of output units.
    """
    if arity < 2:
        raise ValueError("The arity should be at least 2")
    super().__init__(
        num_input_units, num_input_units, arity=arity, semiring=semiring, num_folds=num_folds
    )

forward(x)

Source code in cirkit/backend/torch/layers/inner.py

def forward(self, x: Tensor) -> Tensor:
    return self.semiring.prod(x, dim=1, keepdim=False)  # shape (F, H, B, K) -> (F, B, K).

sample(x)

Source code in cirkit/backend/torch/layers/inner.py

def sample(self, x: Tensor) -> tuple[Tensor, None]:
    # Concatenate samples over disjoint variables through a sum
    # x: (F, H, C, K, num_samples, D)
    x = torch.sum(x, dim=1)  # (F, C, K, num_samples, D)
    return x, None

TorchInnerLayer

Bases: TorchLayer, ABC

The abstract base class for inner layers, i.e., either sum or product layers.

Source code in cirkit/backend/torch/layers/inner.py

class TorchInnerLayer(TorchLayer, ABC):
    """The abstract base class for inner layers, i.e., either sum or product layers."""

    def __init__(
        self,
        num_input_units: int,
        num_output_units: int,
        arity: int = 2,
        *,
        semiring: Semiring | None = None,
        num_folds: int = 1,
    ):
        """Initialize an inner layer.

        Args:
            num_input_units: The number of input units.
            num_output_units: The number of output units.
            arity: The arity of the layer.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
            num_folds: The number of channels.
        """
        super().__init__(
            num_input_units, num_output_units, arity=arity, semiring=semiring, num_folds=num_folds
        )

    @property
    def fold_settings(self) -> tuple[Any, ...]:
        pshapes = [(n, p.shape) for n, p in self.params.items()]
        return *self.config.items(), *pshapes

    def __call__(self, x: Tensor) -> Tensor:
        # IGNORE: Idiom for nn.Module.__call__.
        return super().__call__(x)  # type: ignore[no-any-return,misc]

    @abstractmethod
    def forward(self, x: Tensor) -> Tensor:
        """Invoke the forward function.

        Args:
            x: The tensor input to this layer, having shape $(F, H, B, K_i)$, where $F$
                is the number of folds, $H$ is the arity, $B$ is the batch size, and
                $K_i$ is the number of input units.

        Returns:
            Tensor: The tensor output of this layer, having shape $(F, B, K_o)$, where $K_o$
                is the number of output units.
        """

    def sample(self, x: Tensor) -> tuple[Tensor, Tensor | None]:
        """Perform a forward sampling step.

        Args:
            x: A tensor representing the input variable assignments, having shape
                $(F, H, C, K, N, D)$, where $F$ is the number of folds, $H$ is the arity,
                $C$ is the number of channels, $K$ is the numbe rof input units, $N$ is the number
                of samples, $D$ is the number of variables.

        Returns:
            Tensor: A new tensor representing the new variable assignements the layers gives
                as output.

        Raises:
            TypeError: If sampling is not supported by the layer.
        """
        raise TypeError(f"Sampling not implemented for {type(self)}")

fold_settings property

__call__(x)

Source code in cirkit/backend/torch/layers/inner.py

def __call__(self, x: Tensor) -> Tensor:
    # IGNORE: Idiom for nn.Module.__call__.
    return super().__call__(x)  # type: ignore[no-any-return,misc]

__init__(num_input_units, num_output_units, arity=2, *, semiring=None, num_folds=1)

Initialize an inner layer.

Parameters:

Name	Type	Description	Default
num_input_units	int	The number of input units.	required
num_output_units	int	The number of output units.	required
arity	int	The arity of the layer.	2
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None
num_folds	int	The number of channels.	1

Source code in cirkit/backend/torch/layers/inner.py

def __init__(
    self,
    num_input_units: int,
    num_output_units: int,
    arity: int = 2,
    *,
    semiring: Semiring | None = None,
    num_folds: int = 1,
):
    """Initialize an inner layer.

    Args:
        num_input_units: The number of input units.
        num_output_units: The number of output units.
        arity: The arity of the layer.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
        num_folds: The number of channels.
    """
    super().__init__(
        num_input_units, num_output_units, arity=arity, semiring=semiring, num_folds=num_folds
    )

forward(x) abstractmethod

Invoke the forward function.

Parameters:

Name	Type	Description	Default
x	Tensor	The tensor input to this layer, having shape \((F, H, B, K_i)\), where \(F\) is the number of folds, \(H\) is the arity, \(B\) is the batch size, and \(K_i\) is the number of input units.	required

Returns:

Name	Type	Description
Tensor	Tensor	The tensor output of this layer, having shape \((F, B, K_o)\), where \(K_o\) is the number of output units.

Source code in cirkit/backend/torch/layers/inner.py

@abstractmethod
def forward(self, x: Tensor) -> Tensor:
    """Invoke the forward function.

    Args:
        x: The tensor input to this layer, having shape $(F, H, B, K_i)$, where $F$
            is the number of folds, $H$ is the arity, $B$ is the batch size, and
            $K_i$ is the number of input units.

    Returns:
        Tensor: The tensor output of this layer, having shape $(F, B, K_o)$, where $K_o$
            is the number of output units.
    """

sample(x)

Perform a forward sampling step.

Parameters:

Name	Type	Description	Default
x	Tensor	A tensor representing the input variable assignments, having shape \((F, H, C, K, N, D)\), where \(F\) is the number of folds, \(H\) is the arity, \(C\) is the number of channels, \(K\) is the numbe rof input units, \(N\) is the number of samples, \(D\) is the number of variables.	required

Returns:

Name	Type	Description
Tensor	tuple[Tensor, Tensor | None]	A new tensor representing the new variable assignements the layers gives as output.

Raises:

Type	Description
TypeError	If sampling is not supported by the layer.

Source code in cirkit/backend/torch/layers/inner.py

def sample(self, x: Tensor) -> tuple[Tensor, Tensor | None]:
    """Perform a forward sampling step.

    Args:
        x: A tensor representing the input variable assignments, having shape
            $(F, H, C, K, N, D)$, where $F$ is the number of folds, $H$ is the arity,
            $C$ is the number of channels, $K$ is the numbe rof input units, $N$ is the number
            of samples, $D$ is the number of variables.

    Returns:
        Tensor: A new tensor representing the new variable assignements the layers gives
            as output.

    Raises:
        TypeError: If sampling is not supported by the layer.
    """
    raise TypeError(f"Sampling not implemented for {type(self)}")

TorchInputLayer

Bases: TorchLayer, ABC

The abstract base class for torch input layers.

Source code in cirkit/backend/torch/layers/input.py

class TorchInputLayer(TorchLayer, ABC):
    """The abstract base class for torch input layers."""

    def __init__(
        self,
        scope_idx: Tensor,
        num_output_units: int,
        *,
        num_channels: int = 1,
        semiring: Semiring | None = None,
    ) -> None:
        r"""Initialize a torch input layer.

        Args:
            scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
                $D$ is the number of variables on which the input layers in each fold are defined
                on. Alternatively, a tensor of shape $(D,)$ can be specified, which will be
                interpreted as a tensor of shape $(1, D)$, i.e., with $F = 1$.
            num_output_units: The number of output units.
            num_channels: The number of channels.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

        Raises:
            ValueError: If the scope index is not a vector or a matrix.
        """
        if len(scope_idx.shape) == 1:
            scope_idx = scope_idx.unsqueeze(dim=0)
        elif len(scope_idx.shape) > 2:
            raise ValueError(f"The scope index must be a matrix, but found shape {scope_idx.shape}")
        num_folds, num_variables = scope_idx.shape
        super().__init__(
            num_variables,
            num_output_units,
            arity=num_channels,
            num_folds=num_folds,
            semiring=semiring,
        )
        self.register_buffer("_scope_idx", scope_idx)

    @property
    def scope_idx(self) -> Tensor:
        """Retrieve the scope index tensor.

        Returns:
            The scope index tensor.
        """
        return self._scope_idx

    @property
    def num_variables(self) -> int:
        """The number of variables the input layer is defined on.

        Returns:
            The number of variables.
        """
        return self.num_input_units

    @property
    def num_channels(self) -> int:
        """The number of channels per variable.

        Returns:
            The number of channels.
        """
        return self.arity

    @property
    @abstractmethod
    def config(self) -> Mapping[str, Any]:
        ...

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        return {}

    @property
    def fold_settings(self) -> tuple[Any, ...]:
        pshapes = [(n, p.shape) for n, p in self.params.items()]
        return self.num_variables, *self.config.items(), *pshapes

    def integrate(self) -> Tensor:
        r"""Integrate an input layer over all its variables' domain.

        Returns:
            Tensor: The tensor result of the integration, having shape $(F, K)$, where
                $F$ is the number of folds and $K$ is the number of output units.

        Raises:
            TypeError: If integration is not supported by the layer.
        """
        raise TypeError(f"Integration is not supported for layers of type {type(self)}")

    def sample(self, num_samples: int = 1) -> Tensor:
        r"""If the input layer encodes a probability distribution, then sample from it.

        Args:
            num_samples: The number of data points to sample.

        Returns:
            Tensor: The tensorized sample, having shape $(F, C, K, N)$, where
                $F$ is the number of folds, $K$ is the number of output units,
                $C$ is the number of channels, and $N$ is the number of samples.

        Raises:
            TypeError: If sampling is not supported by the layer.
        """
        raise TypeError(f"Sampling is not supported for layers of type {type(self)}")

    def extra_repr(self) -> str:
        return (
            "  ".join(
                [
                    f"folds: {self.num_folds}",
                    f"channels: {self.num_channels}",
                    f"variables: {self.num_variables}",
                    f"output-units: {self.num_output_units}",
                ]
            )
            + "\n"
            + f"input-shape: {(self.num_folds, self.arity, -1, self.num_input_units)}"
            + "\n"
            + f"output-shape: {(self.num_folds, -1, self.num_output_units)}"
        )

config abstractmethod property

fold_settings property

num_channels property

The number of channels per variable.

Returns:

Type	Description
int	The number of channels.

num_variables property

The number of variables the input layer is defined on.

Returns:

Type	Description
int	The number of variables.

params property

scope_idx property

Retrieve the scope index tensor.

Returns:

Type	Description
Tensor	The scope index tensor.

__init__(scope_idx, num_output_units, *, num_channels=1, semiring=None)

Initialize a torch input layer.

Parameters:

Name	Type	Description	Default
scope_idx	Tensor	A tensor of shape \((F, D)\), where \(F\) is the number of folds, and \(D\) is the number of variables on which the input layers in each fold are defined on. Alternatively, a tensor of shape \((D,)\) can be specified, which will be interpreted as a tensor of shape \((1, D)\), i.e., with \(F = 1\).	required
num_output_units	int	The number of output units.	required
num_channels	int	The number of channels.	1
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None

Raises:

Type	Description
ValueError	If the scope index is not a vector or a matrix.

Source code in cirkit/backend/torch/layers/input.py

def __init__(
    self,
    scope_idx: Tensor,
    num_output_units: int,
    *,
    num_channels: int = 1,
    semiring: Semiring | None = None,
) -> None:
    r"""Initialize a torch input layer.

    Args:
        scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
            $D$ is the number of variables on which the input layers in each fold are defined
            on. Alternatively, a tensor of shape $(D,)$ can be specified, which will be
            interpreted as a tensor of shape $(1, D)$, i.e., with $F = 1$.
        num_output_units: The number of output units.
        num_channels: The number of channels.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

    Raises:
        ValueError: If the scope index is not a vector or a matrix.
    """
    if len(scope_idx.shape) == 1:
        scope_idx = scope_idx.unsqueeze(dim=0)
    elif len(scope_idx.shape) > 2:
        raise ValueError(f"The scope index must be a matrix, but found shape {scope_idx.shape}")
    num_folds, num_variables = scope_idx.shape
    super().__init__(
        num_variables,
        num_output_units,
        arity=num_channels,
        num_folds=num_folds,
        semiring=semiring,
    )
    self.register_buffer("_scope_idx", scope_idx)

extra_repr()

Source code in cirkit/backend/torch/layers/input.py

def extra_repr(self) -> str:
    return (
        "  ".join(
            [
                f"folds: {self.num_folds}",
                f"channels: {self.num_channels}",
                f"variables: {self.num_variables}",
                f"output-units: {self.num_output_units}",
            ]
        )
        + "\n"
        + f"input-shape: {(self.num_folds, self.arity, -1, self.num_input_units)}"
        + "\n"
        + f"output-shape: {(self.num_folds, -1, self.num_output_units)}"
    )

integrate()

Integrate an input layer over all its variables' domain.

Returns:

Name	Type	Description
Tensor	Tensor	The tensor result of the integration, having shape \((F, K)\), where \(F\) is the number of folds and \(K\) is the number of output units.

Raises:

Type	Description
TypeError	If integration is not supported by the layer.

Source code in cirkit/backend/torch/layers/input.py

def integrate(self) -> Tensor:
    r"""Integrate an input layer over all its variables' domain.

    Returns:
        Tensor: The tensor result of the integration, having shape $(F, K)$, where
            $F$ is the number of folds and $K$ is the number of output units.

    Raises:
        TypeError: If integration is not supported by the layer.
    """
    raise TypeError(f"Integration is not supported for layers of type {type(self)}")

sample(num_samples=1)

If the input layer encodes a probability distribution, then sample from it.

Parameters:

Name	Type	Description	Default
num_samples	int	The number of data points to sample.	1

Returns:

Name	Type	Description
Tensor	Tensor	The tensorized sample, having shape \((F, C, K, N)\), where \(F\) is the number of folds, \(K\) is the number of output units, \(C\) is the number of channels, and \(N\) is the number of samples.

Raises:

Type	Description
TypeError	If sampling is not supported by the layer.

Source code in cirkit/backend/torch/layers/input.py

def sample(self, num_samples: int = 1) -> Tensor:
    r"""If the input layer encodes a probability distribution, then sample from it.

    Args:
        num_samples: The number of data points to sample.

    Returns:
        Tensor: The tensorized sample, having shape $(F, C, K, N)$, where
            $F$ is the number of folds, $K$ is the number of output units,
            $C$ is the number of channels, and $N$ is the number of samples.

    Raises:
        TypeError: If sampling is not supported by the layer.
    """
    raise TypeError(f"Sampling is not supported for layers of type {type(self)}")

TorchKroneckerLayer

Bases: TorchInnerLayer

The Kronecker product layer, which computes the Kronecker product of the input vectors it receives as input. See the symbolic KroneckerLayer for more details.

Source code in cirkit/backend/torch/layers/inner.py

class TorchKroneckerLayer(TorchInnerLayer):
    """The Kronecker product layer, which computes the Kronecker product of the input vectors
    it receives as input.
    See the symbolic [KroneckerLayer][cirkit.symbolic.layers.KroneckerLayer] for more details.
    """

    def __init__(
        self,
        num_input_units: int,
        arity: int = 2,
        *,
        semiring: Semiring | None = None,
        num_folds: int = 1,
    ):
        """Initialize a Kronecker product layer.

        Args:
            num_input_units: The number of input units. The number of output units is the power of
                the number of input units to the arity.
            arity: The arity of the layer. Defaults to 2 (which is the only supported arity).
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
            num_folds: The number of channels.

        Raises:
            NotImplementedError: If the arity is not 2.
            ValueError: If the number of input units is not the same as the number of output units.
        """
        # TODO: generalize kronecker layer as to support a greater arity
        if arity != 2:
            raise NotImplementedError("Kronecker only implemented for binary product units.")
        super().__init__(
            num_input_units,
            num_input_units**arity,
            arity=arity,
            semiring=semiring,
            num_folds=num_folds,
        )

    @property
    def config(self) -> Mapping[str, Any]:
        return {
            "num_input_units": self.num_input_units,
            "arity": self.arity,
        }

    def forward(self, x: Tensor) -> Tensor:
        x0 = x[:, 0].unsqueeze(dim=-1)  # shape (F, B, Ki, 1).
        x1 = x[:, 1].unsqueeze(dim=-2)  # shape (F, B, 1, Ki).
        # shape (F, B, Ki, Ki) -> (F, B, Ko=Ki**2).
        return self.semiring.mul(x0, x1).flatten(start_dim=-2)

    def sample(self, x: Tensor) -> tuple[Tensor, Tensor | None]:
        # x: (F, H, C, K, num_samples, D)
        x0 = x[:, 0].unsqueeze(dim=3)  # (F, C, Ki, 1, num_samples, D)
        x1 = x[:, 1].unsqueeze(dim=2)  # (F, C, 1, Ki, num_samples, D)
        # shape (F, C, Ki, Ki, num_samples, D) -> (F, C, Ko=Ki**2, num_samples, D)
        x = x0 + x1
        return torch.flatten(x, start_dim=2, end_dim=3), None

config property

__init__(num_input_units, arity=2, *, semiring=None, num_folds=1)

Initialize a Kronecker product layer.

Parameters:

Name	Type	Description	Default
num_input_units	int	The number of input units. The number of output units is the power of the number of input units to the arity.	required
arity	int	The arity of the layer. Defaults to 2 (which is the only supported arity).	2
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None
num_folds	int	The number of channels.	1

Raises:

Type	Description
NotImplementedError	If the arity is not 2.
ValueError	If the number of input units is not the same as the number of output units.

Source code in cirkit/backend/torch/layers/inner.py

def __init__(
    self,
    num_input_units: int,
    arity: int = 2,
    *,
    semiring: Semiring | None = None,
    num_folds: int = 1,
):
    """Initialize a Kronecker product layer.

    Args:
        num_input_units: The number of input units. The number of output units is the power of
            the number of input units to the arity.
        arity: The arity of the layer. Defaults to 2 (which is the only supported arity).
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
        num_folds: The number of channels.

    Raises:
        NotImplementedError: If the arity is not 2.
        ValueError: If the number of input units is not the same as the number of output units.
    """
    # TODO: generalize kronecker layer as to support a greater arity
    if arity != 2:
        raise NotImplementedError("Kronecker only implemented for binary product units.")
    super().__init__(
        num_input_units,
        num_input_units**arity,
        arity=arity,
        semiring=semiring,
        num_folds=num_folds,
    )

forward(x)

Source code in cirkit/backend/torch/layers/inner.py

def forward(self, x: Tensor) -> Tensor:
    x0 = x[:, 0].unsqueeze(dim=-1)  # shape (F, B, Ki, 1).
    x1 = x[:, 1].unsqueeze(dim=-2)  # shape (F, B, 1, Ki).
    # shape (F, B, Ki, Ki) -> (F, B, Ko=Ki**2).
    return self.semiring.mul(x0, x1).flatten(start_dim=-2)

sample(x)

Source code in cirkit/backend/torch/layers/inner.py

def sample(self, x: Tensor) -> tuple[Tensor, Tensor | None]:
    # x: (F, H, C, K, num_samples, D)
    x0 = x[:, 0].unsqueeze(dim=3)  # (F, C, Ki, 1, num_samples, D)
    x1 = x[:, 1].unsqueeze(dim=2)  # (F, C, 1, Ki, num_samples, D)
    # shape (F, C, Ki, Ki, num_samples, D) -> (F, C, Ko=Ki**2, num_samples, D)
    x = x0 + x1
    return torch.flatten(x, start_dim=2, end_dim=3), None

TorchLayer

Bases: AbstractTorchModule, ABC

The abstract base class for all layers implemented in torch.

Source code in cirkit/backend/torch/layers/base.py

class TorchLayer(AbstractTorchModule, ABC):
    """The abstract base class for all layers implemented in torch."""

    def __init__(
        self,
        num_input_units: int,
        num_output_units: int,
        arity: int = 1,
        *,
        semiring: Semiring | None = None,
        num_folds: int = 1,
    ) -> None:
        """Initialize a layer.

        Args:
            num_input_units: The number of input units.
            num_output_units: The number of output units.
            arity: The arity of the layer.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
            num_folds: The number of folds.

        Raises:
            ValueError: If the number of input units is negative.
            ValueError: If the number of output units is not positive.
            VAlueError: If the arity is not positive.
        """
        if num_input_units < 0:
            raise ValueError("The number of input units must be non-negative")
        if num_output_units <= 0:
            raise ValueError("The number of output units must be positive")
        if arity <= 0:
            raise ValueError("The arity must be positive")
        super().__init__(num_folds=num_folds)
        self.num_input_units = num_input_units
        self.num_output_units = num_output_units
        self.arity = arity
        self.semiring = semiring if semiring is not None else SumProductSemiring

    @property
    @abstractmethod
    def config(self) -> Mapping[str, Any]:
        """Retrieves the configuration of the layer, i.e., a dictionary mapping hyperparameters
        of the layer to their values. The hyperparameter names must match the argument names in
        the ```__init__``` method.

        Returns:
            Mapping[str, Any]: A dictionary from hyperparameter names to their value.
        """

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        """Retrieve the torch parameters of the layer, i.e., a dictionary mapping the names of
        the torch parameters to the actual torch parameter instance. The parameter names must
        match the argument names in the```__init__``` method.

        Returns:
            Mapping[str, TorchParameter]: A dictionary from parameter names to the corresponding
                torch parameter instance.
        """
        return {}

    @property
    def sub_modules(self) -> Mapping[str, "TorchLayer"]:
        """Retrieve a dictionary mapping string identifiers to torch sub-module layers.,
        that must be passed to the ```__init__``` method of the top-level layer

        Returns:
            A dictionary of torch modules.
        """
        return {}

    @cached_property
    def num_parameters(self) -> int:
        """Retrieve the number of scalar parameters. Note that if a parameter is complex-valued,
        this will double count them.

        Returns:
            The number of scalar parameters.
        """
        return sum(2 * p.numel() if torch.is_complex(p) else p.numel() for p in self.parameters())

    @cached_property
    def num_buffers(self) -> int:
        """Retrieve the number of scalar buffers. Note that if a buffer is complex-valued,
        this will double count them.

        Returns:
            The number of scalar buffers.
        """
        return sum(2 * b.numel() if torch.is_complex(b) else b.numel() for b in self.buffers())

    def extra_repr(self) -> str:
        return (
            "  ".join(
                [
                    f"folds: {self.num_folds}",
                    f"arity: {self.arity}",
                    f"input-units: {self.num_input_units}",
                    f"output-units: {self.num_output_units}",
                ]
            )
            + "\n"
            + f"input-shape: {(self.num_folds, self.arity, -1, self.num_input_units)}"
            + "\n"
            + f"output-shape: {(self.num_folds, -1, self.num_output_units)}"
        )

arity = arity instance-attribute

config abstractmethod property

Retrieves the configuration of the layer, i.e., a dictionary mapping hyperparameters of the layer to their values. The hyperparameter names must match the argument names in the __init__ method.

Returns:

Type	Description
Mapping[str, Any]	Mapping[str, Any]: A dictionary from hyperparameter names to their value.

num_buffers cached property

Retrieve the number of scalar buffers. Note that if a buffer is complex-valued, this will double count them.

Returns:

Type	Description
int	The number of scalar buffers.

num_input_units = num_input_units instance-attribute

num_output_units = num_output_units instance-attribute

num_parameters cached property

Retrieve the number of scalar parameters. Note that if a parameter is complex-valued, this will double count them.

Returns:

Type	Description
int	The number of scalar parameters.

params property

Retrieve the torch parameters of the layer, i.e., a dictionary mapping the names of the torch parameters to the actual torch parameter instance. The parameter names must match the argument names in the__init__ method.

Returns:

Type	Description
Mapping[str, TorchParameter]	Mapping[str, TorchParameter]: A dictionary from parameter names to the corresponding torch parameter instance.

semiring = semiring if semiring is not None else SumProductSemiring instance-attribute

sub_modules property

Retrieve a dictionary mapping string identifiers to torch sub-module layers., that must be passed to the __init__ method of the top-level layer

Returns:

Type	Description
Mapping[str, TorchLayer]	A dictionary of torch modules.

__init__(num_input_units, num_output_units, arity=1, *, semiring=None, num_folds=1)

Initialize a layer.

Parameters:

Name	Type	Description	Default
num_input_units	int	The number of input units.	required
num_output_units	int	The number of output units.	required
arity	int	The arity of the layer.	1
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None
num_folds	int	The number of folds.	1

Raises:

Type	Description
ValueError	If the number of input units is negative.
ValueError	If the number of output units is not positive.
VAlueError	If the arity is not positive.

Source code in cirkit/backend/torch/layers/base.py

def __init__(
    self,
    num_input_units: int,
    num_output_units: int,
    arity: int = 1,
    *,
    semiring: Semiring | None = None,
    num_folds: int = 1,
) -> None:
    """Initialize a layer.

    Args:
        num_input_units: The number of input units.
        num_output_units: The number of output units.
        arity: The arity of the layer.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
        num_folds: The number of folds.

    Raises:
        ValueError: If the number of input units is negative.
        ValueError: If the number of output units is not positive.
        VAlueError: If the arity is not positive.
    """
    if num_input_units < 0:
        raise ValueError("The number of input units must be non-negative")
    if num_output_units <= 0:
        raise ValueError("The number of output units must be positive")
    if arity <= 0:
        raise ValueError("The arity must be positive")
    super().__init__(num_folds=num_folds)
    self.num_input_units = num_input_units
    self.num_output_units = num_output_units
    self.arity = arity
    self.semiring = semiring if semiring is not None else SumProductSemiring

extra_repr()

Source code in cirkit/backend/torch/layers/base.py

def extra_repr(self) -> str:
    return (
        "  ".join(
            [
                f"folds: {self.num_folds}",
                f"arity: {self.arity}",
                f"input-units: {self.num_input_units}",
                f"output-units: {self.num_output_units}",
            ]
        )
        + "\n"
        + f"input-shape: {(self.num_folds, self.arity, -1, self.num_input_units)}"
        + "\n"
        + f"output-shape: {(self.num_folds, -1, self.num_output_units)}"
    )

TorchLogPartitionLayer

Bases: TorchConstantLayer

An input layer having empty scope and computing a constant value.

Source code in cirkit/backend/torch/layers/input.py

class TorchConstantValueLayer(TorchConstantLayer):
    """An input layer having empty scope and computing a constant value."""

    def __init__(
        self,
        num_output_units: int,
        *,
        log_space: bool = False,
        value: TorchParameter,
        semiring: Semiring | None = None,
    ) -> None:
        r"""Initialize a constant value input layer.

        Args:
            num_output_units: The number of output units.
            log_space: Whether the given value is in the log-space, i.e., this constant
                layer should encode functions $\exp(x)$ rather than just x.
            value: The tensor value encoded by the layer, given by a parameter of shape $(F, K)$,
                where $F$ is the number of folds and $K$ is the numer of output units.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

        Raises:
            ValueError: If the number of folds of the shape of the given value is incorrect.
        """
        super().__init__(
            num_output_units,
            value.num_folds,
            semiring=semiring,
        )
        if value.num_folds != self.num_folds:
            raise ValueError(
                f"The value must have number of folds {self.num_folds}, "
                f"but found {value.num_folds}"
            )
        if value.shape != (num_output_units,):
            raise ValueError(
                f"The shape of the value must be ({num_output_units},), " f"but found {value.shape}"
            )
        self.value = value
        self.log_space = log_space
        self._source_semiring = LSESumSemiring if log_space else SumProductSemiring

    @property
    def config(self) -> Mapping[str, Any]:
        return {"num_output_units": self.num_output_units, "log_space": self.log_space}

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        return {"value": self.value}

    def forward(self, batch_size: int) -> Tensor:
        value = self.value()  # (F, Ko)
        # value: (F, B, Ko)
        value = value.unsqueeze(dim=1).expand(value.shape[0], batch_size, value.shape[1])
        return self.semiring.map_from(value, self._source_semiring)

_source_semiring = LSESumSemiring if log_space else SumProductSemiring instance-attribute

config property

log_space = log_space instance-attribute

params property

value = value instance-attribute

__init__(num_output_units, *, log_space=False, value, semiring=None)

Initialize a constant value input layer.

Parameters:

Name	Type	Description	Default
num_output_units	int	The number of output units.	required
log_space	bool	Whether the given value is in the log-space, i.e., this constant layer should encode functions \(\exp(x)\) rather than just x.	False
value	TorchParameter	The tensor value encoded by the layer, given by a parameter of shape \((F, K)\), where \(F\) is the number of folds and \(K\) is the numer of output units.	required
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None

Raises:

Type	Description
ValueError	If the number of folds of the shape of the given value is incorrect.

Source code in cirkit/backend/torch/layers/input.py

def __init__(
    self,
    num_output_units: int,
    *,
    log_space: bool = False,
    value: TorchParameter,
    semiring: Semiring | None = None,
) -> None:
    r"""Initialize a constant value input layer.

    Args:
        num_output_units: The number of output units.
        log_space: Whether the given value is in the log-space, i.e., this constant
            layer should encode functions $\exp(x)$ rather than just x.
        value: The tensor value encoded by the layer, given by a parameter of shape $(F, K)$,
            where $F$ is the number of folds and $K$ is the numer of output units.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].

    Raises:
        ValueError: If the number of folds of the shape of the given value is incorrect.
    """
    super().__init__(
        num_output_units,
        value.num_folds,
        semiring=semiring,
    )
    if value.num_folds != self.num_folds:
        raise ValueError(
            f"The value must have number of folds {self.num_folds}, "
            f"but found {value.num_folds}"
        )
    if value.shape != (num_output_units,):
        raise ValueError(
            f"The shape of the value must be ({num_output_units},), " f"but found {value.shape}"
        )
    self.value = value
    self.log_space = log_space
    self._source_semiring = LSESumSemiring if log_space else SumProductSemiring

forward(batch_size)

Source code in cirkit/backend/torch/layers/input.py

def forward(self, batch_size: int) -> Tensor:
    value = self.value()  # (F, Ko)
    # value: (F, B, Ko)
    value = value.unsqueeze(dim=1).expand(value.shape[0], batch_size, value.shape[1])
    return self.semiring.map_from(value, self._source_semiring)

TorchPolynomialLayer

Bases: TorchInputFunctionLayer

The polynomial input layer, evaluating a vector of parameterized polynomials.

Source code in cirkit/backend/torch/layers/input.py

class TorchPolynomialLayer(TorchInputFunctionLayer):
    """The polynomial input layer, evaluating a vector of parameterized polynomials."""

    def __init__(
        self,
        scope_idx: Tensor,
        num_output_units: int,
        num_channels: int = 1,
        *,
        degree: int,
        coeff: TorchParameter,
        semiring: Semiring | None = None,
    ) -> None:
        r"""Initialize a polynomial layer.

        Args:
            scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
                $D$ is the number of variables on which the input layers in each fold are defined
                on. Alternatively, a tensor of shape $(D,)$ can be specified, which will be
                interpreted as a tensor of shape $(1, D)$, i.e., with $F = 1$.
            num_output_units: The number of output units.
            num_channels: The number of channels.
            degree: The degree of polynomial.
            coeff: The coefficient parameter, having shape $(F, K, \mathsf{degree} + 1)$, where $K$ is the number
                of output units.

        Raises:
            ValueError: If the scope contains more than one variable.
            ValueError: If the coefficients is not correct.
        """
        num_variables = scope_idx.shape[-1]
        if num_variables != 1:
            raise ValueError("The Polynomial layer encodes a univariate distribution")
        if num_channels != 1:
            raise ValueError("The Polynomial layer encodes a univariate distribution")
        super().__init__(
            scope_idx,
            num_output_units,
            num_channels=num_channels,
            semiring=semiring,
        )
        self.degree = degree
        if not self._valid_parameters_shape(coeff):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._coeff_shape} for 'coeff', found"
                f"{coeff.num_folds} and {coeff.shape}, respectively"
            )
        self.coeff = coeff

    def _valid_parameters_shape(self, p: TorchParameter) -> bool:
        if p.num_folds != self.num_folds:
            return False
        return p.shape == self._coeff_shape

    @property
    def _coeff_shape(self) -> tuple[int, ...]:
        return self.num_output_units, self.degree + 1

    @staticmethod
    def _polyval(coeff: Tensor, x: Tensor) -> Tensor:
        r"""Evaluate polynomial given coefficients and point, with the shape for PolynomialLayer.

        Args:
            coeff: The coefficients of the polynomial, shape $(F, K_o, \mathsf{degree} + 1)$.
            x: The point of the variable, shape $(F, H, B, K_i)$, where $H=K_i=1$.

        Returns:
            Tensor: The value of the polymonial, shape $(F, B, K_o)$.
        """
        x = x.squeeze(dim=1)  # shape (F, H=1, B, Ki=1) -> (F, B, 1).
        y = x.new_zeros(*x.shape[:-1], coeff.shape[-2])  # shape (F, B, Ko).

        # TODO: iterating over dim=2 is inefficient
        for a_n in reversed(
            coeff.unbind(dim=2)
        ):  # Reverse iterator of the degree axis, shape (F, Ko).
            # a_n shape (F, Ko) -> (F, 1, Ko).
            y = torch.addcmul(a_n.unsqueeze(dim=1), x, y)  # y = a_n + x * y, by Horner's method.
        return y  # shape (F, B, Ko).

    @property
    def config(self) -> Mapping[str, Any]:
        return {
            "num_output_units": self.num_output_units,
            "num_channels": self.num_channels,
            "degree": self.degree,
        }

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        return {"coeff": self.coeff}

    def forward(self, x: Tensor) -> Tensor:
        coeff = self.coeff()  # shape (F, Ko, dp1)
        return self.semiring.map_from(TorchPolynomialLayer._polyval(coeff, x), SumProductSemiring)

_coeff_shape property

coeff = coeff instance-attribute

config property

degree = degree instance-attribute

params property

__init__(scope_idx, num_output_units, num_channels=1, *, degree, coeff, semiring=None)

Initialize a polynomial layer.

Parameters:

Name	Type	Description	Default
scope_idx	Tensor	A tensor of shape \((F, D)\), where \(F\) is the number of folds, and \(D\) is the number of variables on which the input layers in each fold are defined on. Alternatively, a tensor of shape \((D,)\) can be specified, which will be interpreted as a tensor of shape \((1, D)\), i.e., with \(F = 1\).	required
num_output_units	int	The number of output units.	required
num_channels	int	The number of channels.	1
degree	int	The degree of polynomial.	required
coeff	TorchParameter	The coefficient parameter, having shape \((F, K, \mathsf{degree} + 1)\), where \(K\) is the number of output units.	required

Raises:

Type	Description
ValueError	If the scope contains more than one variable.
ValueError	If the coefficients is not correct.

Source code in cirkit/backend/torch/layers/input.py

def __init__(
    self,
    scope_idx: Tensor,
    num_output_units: int,
    num_channels: int = 1,
    *,
    degree: int,
    coeff: TorchParameter,
    semiring: Semiring | None = None,
) -> None:
    r"""Initialize a polynomial layer.

    Args:
        scope_idx: A tensor of shape $(F, D)$, where $F$ is the number of folds, and
            $D$ is the number of variables on which the input layers in each fold are defined
            on. Alternatively, a tensor of shape $(D,)$ can be specified, which will be
            interpreted as a tensor of shape $(1, D)$, i.e., with $F = 1$.
        num_output_units: The number of output units.
        num_channels: The number of channels.
        degree: The degree of polynomial.
        coeff: The coefficient parameter, having shape $(F, K, \mathsf{degree} + 1)$, where $K$ is the number
            of output units.

    Raises:
        ValueError: If the scope contains more than one variable.
        ValueError: If the coefficients is not correct.
    """
    num_variables = scope_idx.shape[-1]
    if num_variables != 1:
        raise ValueError("The Polynomial layer encodes a univariate distribution")
    if num_channels != 1:
        raise ValueError("The Polynomial layer encodes a univariate distribution")
    super().__init__(
        scope_idx,
        num_output_units,
        num_channels=num_channels,
        semiring=semiring,
    )
    self.degree = degree
    if not self._valid_parameters_shape(coeff):
        raise ValueError(
            f"Expected number of folds {self.num_folds} "
            f"and shape {self._coeff_shape} for 'coeff', found"
            f"{coeff.num_folds} and {coeff.shape}, respectively"
        )
    self.coeff = coeff

_polyval(coeff, x) staticmethod

Evaluate polynomial given coefficients and point, with the shape for PolynomialLayer.

Parameters:

Name	Type	Description	Default
coeff	Tensor	The coefficients of the polynomial, shape \((F, K_o, \mathsf{degree} + 1)\).	required
x	Tensor	The point of the variable, shape \((F, H, B, K_i)\), where \(H=K_i=1\).	required

Returns:

Name	Type	Description
Tensor	Tensor	The value of the polymonial, shape \((F, B, K_o)\).

Source code in cirkit/backend/torch/layers/input.py

@staticmethod
def _polyval(coeff: Tensor, x: Tensor) -> Tensor:
    r"""Evaluate polynomial given coefficients and point, with the shape for PolynomialLayer.

    Args:
        coeff: The coefficients of the polynomial, shape $(F, K_o, \mathsf{degree} + 1)$.
        x: The point of the variable, shape $(F, H, B, K_i)$, where $H=K_i=1$.

    Returns:
        Tensor: The value of the polymonial, shape $(F, B, K_o)$.
    """
    x = x.squeeze(dim=1)  # shape (F, H=1, B, Ki=1) -> (F, B, 1).
    y = x.new_zeros(*x.shape[:-1], coeff.shape[-2])  # shape (F, B, Ko).

    # TODO: iterating over dim=2 is inefficient
    for a_n in reversed(
        coeff.unbind(dim=2)
    ):  # Reverse iterator of the degree axis, shape (F, Ko).
        # a_n shape (F, Ko) -> (F, 1, Ko).
        y = torch.addcmul(a_n.unsqueeze(dim=1), x, y)  # y = a_n + x * y, by Horner's method.
    return y  # shape (F, B, Ko).

_valid_parameters_shape(p)

Source code in cirkit/backend/torch/layers/input.py

def _valid_parameters_shape(self, p: TorchParameter) -> bool:
    if p.num_folds != self.num_folds:
        return False
    return p.shape == self._coeff_shape

forward(x)

Source code in cirkit/backend/torch/layers/input.py

def forward(self, x: Tensor) -> Tensor:
    coeff = self.coeff()  # shape (F, Ko, dp1)
    return self.semiring.map_from(TorchPolynomialLayer._polyval(coeff, x), SumProductSemiring)

TorchSumLayer

Bases: TorchInnerLayer

The sum layer torch implementation. See the symbolic SumLayer for more details.

Source code in cirkit/backend/torch/layers/inner.py

class TorchSumLayer(TorchInnerLayer):
    """The sum layer torch implementation.
    See the symbolic [SumLayer][cirkit.symbolic.layers.SumLayer] for more details."""

    def __init__(
        self,
        num_input_units: int,
        num_output_units: int,
        arity: int = 1,
        *,
        weight: TorchParameter,
        semiring: Semiring | None = None,
        num_folds: int = 1,
    ):
        r"""Initialize a sum layer.

        Args:
            num_input_units: The number of input units.
            num_output_units: The number of output units.
            arity: The arity of the layer.
            weight: The weight parameter, which must have shape $(F, K_o, K_i\cdot H)$,
                where $F$ is the number of folds, $K_o$ is the number of output units,
                   $K_i$ is the number of input units, and $H$ is the arity.
            semiring: The evaluation semiring.
                Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
            num_folds: The number of channels.

        Raises:
            ValueError: If the arity is not a positive integer.
            ValueError: If the arity, the number of input and output units are incompatible with the
                shape of the weight parameter.
        """
        if arity < 1:
            raise ValueError("The arity must be a positive integer")
        super().__init__(
            num_input_units, num_output_units, arity=arity, semiring=semiring, num_folds=num_folds
        )
        if not self._valid_weight_shape(weight):
            raise ValueError(
                f"Expected number of folds {self.num_folds} "
                f"and shape {self._weight_shape} for 'weight', found"
                f"{weight.num_folds} and {weight.shape}, respectively"
            )
        self.weight = weight

    def _valid_weight_shape(self, w: TorchParameter) -> bool:
        if w.num_folds != self.num_folds:
            return False
        return w.shape == self._weight_shape

    @property
    def _weight_shape(self) -> tuple[int, ...]:
        return self.num_output_units, self.num_input_units * self.arity

    @property
    def config(self) -> Mapping[str, Any]:
        return {
            "num_input_units": self.num_input_units,
            "num_output_units": self.num_output_units,
            "arity": self.arity,
        }

    @property
    def params(self) -> Mapping[str, TorchParameter]:
        return {"weight": self.weight}

    def forward(self, x: Tensor) -> Tensor:
        # x: (F, H, B, Ki) -> (F, B, H * Ki)
        # weight: (F, Ko, H * Ki)
        x = x.permute(0, 2, 1, 3).flatten(start_dim=2)
        weight = self.weight()
        return self.semiring.einsum(
            "fbi,foi->fbo", inputs=(x,), operands=(weight,), dim=-1, keepdim=True
        )  # shape (F, B, K_o).

    def sample(self, x: Tensor) -> tuple[Tensor, Tensor]:
        weight = self.weight()
        negative = torch.any(weight < 0.0)
        normalized = torch.allclose(torch.sum(weight, dim=-1), torch.ones(1, device=weight.device))
        if negative or not normalized:
            raise TypeError("Sampling in sum layers only works with positive weights summing to 1")

        # x: (F, H, C, Ki, num_samples, D) -> (F, C, H * Ki, num_samples, D)
        x = x.permute(0, 2, 1, 3, 4, 5).flatten(2, 3)
        c = x.shape[1]
        num_samples = x.shape[3]
        d = x.shape[4]

        # mixing_distribution: (F, Ko, H * Ki)
        mixing_distribution = torch.distributions.Categorical(probs=weight)

        # mixing_samples: (num_samples, F, Ko) -> (F, Ko, num_samples)
        mixing_samples = mixing_distribution.sample((num_samples,))
        mixing_samples = E.rearrange(mixing_samples, "n f k -> f k n")

        # mixing_indices: (F, C, Ko, num_samples, D)
        mixing_indices = E.repeat(mixing_samples, "f k n -> f c k n d", c=c, d=d)

        # x: (F, C, Ko, num_samples, D)
        x = torch.gather(x, dim=2, index=mixing_indices)
        return x, mixing_samples

_weight_shape property

config property

params property

weight = weight instance-attribute

__init__(num_input_units, num_output_units, arity=1, *, weight, semiring=None, num_folds=1)

Initialize a sum layer.

Parameters:

Name	Type	Description	Default
num_input_units	int	The number of input units.	required
num_output_units	int	The number of output units.	required
arity	int	The arity of the layer.	1
weight	TorchParameter	The weight parameter, which must have shape \((F, K_o, K_i\cdot H)\), where \(F\) is the number of folds, \(K_o\) is the number of output units, \(K_i\) is the number of input units, and \(H\) is the arity.	required
semiring	Semiring | None	The evaluation semiring. Defaults to SumProductSemiring.	None
num_folds	int	The number of channels.	1

Raises:

Type	Description
ValueError	If the arity is not a positive integer.
ValueError	If the arity, the number of input and output units are incompatible with the shape of the weight parameter.

Source code in cirkit/backend/torch/layers/inner.py

def __init__(
    self,
    num_input_units: int,
    num_output_units: int,
    arity: int = 1,
    *,
    weight: TorchParameter,
    semiring: Semiring | None = None,
    num_folds: int = 1,
):
    r"""Initialize a sum layer.

    Args:
        num_input_units: The number of input units.
        num_output_units: The number of output units.
        arity: The arity of the layer.
        weight: The weight parameter, which must have shape $(F, K_o, K_i\cdot H)$,
            where $F$ is the number of folds, $K_o$ is the number of output units,
               $K_i$ is the number of input units, and $H$ is the arity.
        semiring: The evaluation semiring.
            Defaults to [SumProductSemiring][cirkit.backend.torch.semiring.SumProductSemiring].
        num_folds: The number of channels.

    Raises:
        ValueError: If the arity is not a positive integer.
        ValueError: If the arity, the number of input and output units are incompatible with the
            shape of the weight parameter.
    """
    if arity < 1:
        raise ValueError("The arity must be a positive integer")
    super().__init__(
        num_input_units, num_output_units, arity=arity, semiring=semiring, num_folds=num_folds
    )
    if not self._valid_weight_shape(weight):
        raise ValueError(
            f"Expected number of folds {self.num_folds} "
            f"and shape {self._weight_shape} for 'weight', found"
            f"{weight.num_folds} and {weight.shape}, respectively"
        )
    self.weight = weight

_valid_weight_shape(w)

Source code in cirkit/backend/torch/layers/inner.py

def _valid_weight_shape(self, w: TorchParameter) -> bool:
    if w.num_folds != self.num_folds:
        return False
    return w.shape == self._weight_shape

forward(x)

Source code in cirkit/backend/torch/layers/inner.py

def forward(self, x: Tensor) -> Tensor:
    # x: (F, H, B, Ki) -> (F, B, H * Ki)
    # weight: (F, Ko, H * Ki)
    x = x.permute(0, 2, 1, 3).flatten(start_dim=2)
    weight = self.weight()
    return self.semiring.einsum(
        "fbi,foi->fbo", inputs=(x,), operands=(weight,), dim=-1, keepdim=True
    )  # shape (F, B, K_o).

sample(x)

Source code in cirkit/backend/torch/layers/inner.py

def sample(self, x: Tensor) -> tuple[Tensor, Tensor]:
    weight = self.weight()
    negative = torch.any(weight < 0.0)
    normalized = torch.allclose(torch.sum(weight, dim=-1), torch.ones(1, device=weight.device))
    if negative or not normalized:
        raise TypeError("Sampling in sum layers only works with positive weights summing to 1")

    # x: (F, H, C, Ki, num_samples, D) -> (F, C, H * Ki, num_samples, D)
    x = x.permute(0, 2, 1, 3, 4, 5).flatten(2, 3)
    c = x.shape[1]
    num_samples = x.shape[3]
    d = x.shape[4]

    # mixing_distribution: (F, Ko, H * Ki)
    mixing_distribution = torch.distributions.Categorical(probs=weight)

    # mixing_samples: (num_samples, F, Ko) -> (F, Ko, num_samples)
    mixing_samples = mixing_distribution.sample((num_samples,))
    mixing_samples = E.rearrange(mixing_samples, "n f k -> f k n")

    # mixing_indices: (F, C, Ko, num_samples, D)
    mixing_indices = E.repeat(mixing_samples, "f k n -> f c k n d", c=c, d=d)

    # x: (F, C, Ko, num_samples, D)
    x = torch.gather(x, dim=2, index=mixing_indices)
    return x, mixing_samples

pc2qpc(pc, integration_method, net_dim=128, bias=True, input_sharing='f', inner_sharing='c', ff_dim=None, ff_sigma=1.0, learn_ff=False, freeze_mixing_layers=True)

Source code in cirkit/backend/torch/parameters/pic.py

@torch.no_grad()
def pc2qpc(
    pc: TorchCircuit,
    integration_method: str,
    net_dim: int | None = 128,
    bias: bool | None = True,
    input_sharing: str | None = "f",
    inner_sharing: str | None = "c",
    ff_dim: int | None = None,
    ff_sigma: float | None = 1.0,
    learn_ff: bool | None = False,
    freeze_mixing_layers: bool = True,
):
    def param_to_buffer(model: torch.nn.Module):
        """Turns all parameters of a module into buffers."""
        modules = model.modules()
        module = next(modules)
        for name, param in module.named_parameters(recurse=False):
            delattr(module, name)  # Unregister parameter
            module.register_buffer(name, param.data)
        for module in modules:
            param_to_buffer(module)

    qpc = pc  # copy.deepcopy(pc)
    param_to_buffer(qpc)

    for node in qpc.nodes:
        if isinstance(node, TorchCategoricalLayer):
            z_quad = zw_quadrature(
                integration_method=integration_method, nip=node.num_output_units
            )[0]
            if node.logits is None:
                tensor_parameter = node.probs.nodes[0]
                reparam = node.probs.nodes[1] if len(node.probs.nodes) == 2 else None
            else:
                tensor_parameter = node.logits.nodes[0]
                reparam = node.logits.nodes[1] if len(node.logits.nodes) == 2 else None
            input_net = PICInputNet(
                num_variables=node.num_variables * node.num_folds,
                num_param=node.num_categories,
                num_channels=node.num_channels,
                net_dim=net_dim,
                bias=bias,
                sharing=input_sharing,
                ff_dim=ff_dim,
                ff_sigma=ff_sigma,
                learn_ff=learn_ff,
                z_quad=z_quad,
                tensor_parameter=tensor_parameter,
                reparam=reparam,
            )
            if node.logits is None:
                node.probs = input_net
            else:
                node.logits = input_net
        elif isinstance(node, TorchGaussianLayer):
            assert len(node.mean.nodes) <= 2 and len(node.stddev) <= 2
            z_quad = zw_quadrature(
                integration_method=integration_method, nip=node.num_output_units
            )[0]
            node.mean = PICInputNet(
                num_variables=node.num_variables * node.num_folds,
                num_param=1,
                num_channels=node.num_channels,
                net_dim=net_dim,
                bias=bias,
                sharing=input_sharing,
                ff_dim=ff_dim,
                ff_sigma=ff_sigma,
                learn_ff=learn_ff,
                z_quad=z_quad,
                tensor_parameter=node.mean.nodes[0],
                reparam=None if len(node.mean.nodes) == 1 else node.mean.nodes[0],
            )
            node.stddev = PICInputNet(
                num_variables=node.num_variables * node.num_folds,
                num_param=1,
                num_channels=node.num_channels,
                net_dim=net_dim,
                bias=bias,
                sharing=input_sharing,
                ff_dim=ff_dim,
                ff_sigma=ff_sigma,
                learn_ff=learn_ff,
                z_quad=z_quad,
                tensor_parameter=node.stddev.nodes[0],
                reparam=None if len(node.stddev.nodes) == 1 else node.stddev.nodes[0],
            )
        elif isinstance(node, (TorchSumLayer, TorchTuckerLayer, TorchCPTLayer)):
            if isinstance(node, TorchSumLayer) and node.arity > 1:
                weight_nodes = list(node.weight.topological_ordering())
                assert len(weight_nodes) <= 2, (
                    "Sum layers with arity greater than one must have parameters "
                    "with at most 2 computational nodes"
                )
                tensor_parameter = node.weight.nodes[0]
                if freeze_mixing_layers:
                    tensor_parameter._ptensor.fill_(1 / node.arity)
                    tensor_parameter._ptensor.requires_grad = False
                    continue
            else:
                assert (
                    len(node.weight.nodes) == 1
                ), "You are probably using a reparameterization. Do not do that, QPCs are already normalized!"
                tensor_parameter = node.weight.nodes[0]
            weight_shape = list(tensor_parameter._ptensor.shape)
            squeezed_weight_shape = [weight_shape[0]] + [
                dim_size for dim_size in weight_shape[1:] if dim_size != 1
            ]
            assert (
                sum(
                    [
                        dim_size % min(squeezed_weight_shape[1:])
                        for dim_size in squeezed_weight_shape[1:]
                    ]
                )
                == 0
            ), f"Cannot model a sum layer with shape {weight_shape}!"
            is_tucker = isinstance(node, TorchTuckerLayer)
            nip = int(max(squeezed_weight_shape[1:]) ** (0.5 if is_tucker else 1))
            num_dim = sum(
                [int(np.emath.logn(nip, dim_size)) for dim_size in squeezed_weight_shape[1:]]
            )
            z_quad, w_quad = zw_quadrature(integration_method=integration_method, nip=nip)
            node.weight = PICInnerNet(
                num_dim=num_dim,
                num_funcs=node.num_folds,
                perm_dim=tuple(range(1, num_dim + 1)),
                norm_dim=tuple(range(1, num_dim + 1))[-(2 if is_tucker else 1) :],
                net_dim=net_dim,
                bias=bias,
                sharing=inner_sharing,
                ff_dim=ff_dim,
                ff_sigma=ff_sigma,
                learn_ff=learn_ff,
                z_quad=z_quad,
                w_quad=w_quad,
                tensor_parameter=tensor_parameter,
            )
        elif isinstance(node, (TorchHadamardLayer, TorchKroneckerLayer)):
            pass
        else:
            raise NotImplementedError("Layer %s is not yet handled!" % str(type(node)))

zw_quadrature(integration_method, nip, a=-1, b=1, return_log_weight=False, dtype=torch.float32)

Source code in cirkit/backend/torch/parameters/pic.py

def zw_quadrature(
    integration_method: str,
    nip: int,
    a: float | None = -1,
    b: float | None = 1,
    return_log_weight: bool | None = False,
    dtype: torch.dtype | None = torch.float32,
):
    if integration_method == "leggauss":
        z_quad, w_quad = np.polynomial.legendre.leggauss(nip)
        z_quad = (b - a) * (z_quad + 1) / 2 + a
        w_quad = w_quad * (b - a) / 2
    elif integration_method == "midpoint":
        z_quad = np.linspace(a, b, num=nip + 1)
        z_quad = (z_quad[:-1] + z_quad[1:]) / 2
        w_quad = np.full_like(z_quad, (b - a) / nip)
    elif integration_method == "trapezoidal":
        z_quad = np.linspace(a, b, num=nip)
        w_quad = np.full((nip,), (b - a) / (nip - 1))
        w_quad[0] = w_quad[-1] = 0.5 * (b - a) / (nip - 1)
    elif integration_method == "simpson":
        assert nip % 2 == 1, "Number of integration points must be odd"
        z_quad = np.linspace(a, b, num=nip)
        w_quad = np.concatenate(
            [np.ones(1), np.tile(np.array([4, 2]), nip // 2 - 1), np.array([4, 1])]
        )
        w_quad = ((b - a) / (nip - 1)) / 3 * w_quad
    elif integration_method == "hermgauss":
        # https://en.wikipedia.org/wiki/Gauss%E2%80%93Hermite_quadrature
        z_quad, w_quad = np.polynomial.hermite.hermgauss(nip)
    else:
        raise NotImplementedError("Integration method not implemented.")
    z_quad = torch.tensor(z_quad, dtype=dtype)
    w_quad = torch.tensor(w_quad, dtype=dtype)
    w_quad = w_quad.log() if return_log_weight else w_quad
    return z_quad, w_quad