added backward sampling in a layer-wise manner

cirkit/backend/torch/layers/__init__.py

@@ -1,4 +1,6 @@
 from .base import TorchLayer as TorchLayer
+from .inner import ArityBranch as ArityBranch
+from .inner import BackwardSelection as BackwardSelection
 from .inner import TorchHadamardLayer as TorchHadamardLayer
 from .inner import TorchInnerLayer as TorchInnerLayer
 from .inner import TorchKroneckerLayer as TorchKroneckerLayer

cirkit/backend/torch/layers/inner.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
+from dataclasses import dataclass
 from typing import Any
 
 import einops as E
@@ -11,6 +12,28 @@
 from cirkit.backend.torch.semiring import Semiring
 
 
+# (sample_ids, folds, units) — three parallel 1-D tensors of length P, where P is the number
+# of active paths at a layer during top-down (backward) sampling.
+BackwardSelection = tuple[Tensor, Tensor, Tensor]
+
+
+@dataclass
+class ArityBranch:
+    """The local result of `backward_sample` for one arity slot of an inner layer.
+
+    `sample_ids` and `folds` are subsets of the parent's selection (subset = identity
+    when all paths are active at this slot, e.g. Hadamard/CPT). `units` is the unit
+    each path selects in the child at this slot.
+
+    The driver translates `folds` (parent fold) through `entry.in_fold_idx[0]` to find
+    the child concat-fold, then dispatches across children via `entry.in_module_ids[0]`.
+    """
+
+    sample_ids: Tensor
+    folds: Tensor
+    units: Tensor
+
+
 class TorchInnerLayer(TorchLayer, ABC):
     """The abstract base class for inner layers, i.e., either sum or product layers."""
 
@@ -81,6 +104,23 @@ def sample(self, x: Tensor) -> tuple[Tensor, Tensor | None]:
         """
         raise TypeError(f"Sampling not implemented for {type(self)}")
 
+    def backward_sample(self, selection: BackwardSelection) -> list[ArityBranch]:
+        """Perform a single top-down (backward) sampling step at this layer.
+
+        Args:
+            selection: `(sample_ids, folds, units)` — the active paths reaching this layer
+                from above. All three are 1-D tensors of the same length P.
+
+        Returns:
+            A list of length `arity`; entry `h` describes what to propagate to arity slot `h`.
+            Sum layers return a non-trivial subset per slot (paths whose sampled input falls
+            in slot `h`); Hadamard/CPT broadcast the full selection to every slot.
+
+        Raises:
+            TypeError: If backward sampling is not supported by the layer.
+        """
+        raise TypeError(f"Backward sampling not implemented for {type(self).__name__}")
+
 
 class TorchHadamardLayer(TorchInnerLayer):
     """The Hadamard product layer, which computes an element-wise (or Hadamard) product of
@@ -132,6 +172,11 @@ def sample(self, x: Tensor) -> tuple[Tensor, None]:
         x = torch.sum(x, dim=1)  # (F, C, K, num_samples, D)
         return x, None
 
+    def backward_sample(self, selection: BackwardSelection) -> list[ArityBranch]:
+        # Independence across children: broadcast the selection unchanged to every arity slot.
+        sample_ids, folds, units = selection
+        return [ArityBranch(sample_ids, folds, units) for _ in range(self.arity)]
+
 
 class TorchKroneckerLayer(TorchInnerLayer):
     """The Kronecker product layer, which computes the Kronecker product of the input vectors
@@ -298,3 +343,25 @@ def sample(self, x: Tensor) -> tuple[Tensor, Tensor]:
         # x: (F, Ko, num_samples, D)
         x = torch.gather(x, dim=1, index=mixing_indices)
         return x, mixing_samples
+
+    def backward_sample(self, selection: BackwardSelection) -> list[ArityBranch]:
+        # Sample which of (Ki * arity) inputs each path follows, then split by arity slot.
+        sample_ids, folds, units = selection
+        weight = self.weight()
+        # Shared-parameter layers replicate `weight` across folds — collapse fold index.
+        param_folds = folds % weight.shape[0]
+        selected_weights = weight[param_folds, units]  # (P, Ki*H)
+        input_idx = torch.distributions.Categorical(probs=selected_weights).sample()  # (P,)
+        arity_branch = input_idx // self.num_input_units
+        unit_within = input_idx % self.num_input_units
+
+        branches: list[ArityBranch] = []
+        for h in range(self.arity):
+            mask = arity_branch == h
+            if mask.any():
+                idx = mask.nonzero(as_tuple=True)[0]
+                branches.append(ArityBranch(sample_ids[idx], folds[idx], unit_within[idx]))
+            else:
+                empty = torch.empty(0, dtype=torch.long, device=sample_ids.device)
+                branches.append(ArityBranch(empty, empty, empty))
+        return branches

cirkit/backend/torch/layers/input.py

@@ -107,6 +107,25 @@ def sample(self, num_samples: int = 1) -> Tensor:
         """
         raise TypeError(f"Sampling is not supported for layers of type {type(self)}")
 
+    def backward_sample(
+        self,
+        selection: tuple[Tensor, Tensor, Tensor],
+        samples: Tensor,
+    ) -> None:
+        """Perform a top-down sampling step at this input layer, writing into `samples`.
+
+        Args:
+            selection: `(sample_ids, folds, units)` — the active paths reaching this layer.
+            samples: The `(num_samples, num_variables)` output buffer. The layer writes the
+                sampled value for variable `scope_idx[fold, 0]` at row `sample_ids[i]`.
+
+        Raises:
+            TypeError: If backward sampling is not supported by the layer.
+        """
+        raise TypeError(
+            f"Backward sampling not implemented for input layer {type(self).__name__}"
+        )
+
     def extra_repr(self) -> str:
         return (
             "  ".join(
@@ -433,6 +452,24 @@ def sample(self, num_samples: int = 1) -> Tensor:
         samples = samples.permute(1, 2, 0)  # (F, K, N)
         return samples
 
+    def backward_sample(
+        self,
+        selection: tuple[Tensor, Tensor, Tensor],
+        samples: Tensor,
+    ) -> None:
+        sample_ids, folds, units = selection
+        if self.logits is None:
+            assert self.probs is not None
+            logits = torch.log(self.probs())
+        else:
+            logits = self.logits()
+        # Shared-parameter layers replicate `logits` across folds — collapse fold index.
+        param_folds = folds % logits.shape[0]
+        selected_logits = logits[param_folds, units]  # (P, C)
+        sampled = distributions.Categorical(logits=selected_logits).sample()  # (P,)
+        var_indices = self.scope_idx[folds, 0]
+        samples[sample_ids, var_indices] = sampled
+
 
 class TorchBinomialLayer(TorchExpFamilyLayer):
     """The Binomial distribution layer.

cirkit/backend/torch/layers/optimized.py

@@ -5,7 +5,7 @@
 import torch
 from torch import Tensor
 
-from cirkit.backend.torch.layers.inner import TorchInnerLayer
+from cirkit.backend.torch.layers.inner import ArityBranch, BackwardSelection, TorchInnerLayer
 from cirkit.backend.torch.parameters.parameter import TorchParameter
 from cirkit.backend.torch.semiring import Semiring
 
@@ -201,6 +201,15 @@ def sample(self, x: Tensor) -> tuple[Tensor, Tensor]:
         x = torch.gather(x, dim=1, index=mixing_indices)
         return x, mixing_samples
 
+    def backward_sample(self, selection: BackwardSelection) -> list[ArityBranch]:
+        # CPT = fused product + sum: sample a unit, broadcast it to every arity slot.
+        sample_ids, folds, units = selection
+        weight = self.weight()
+        param_folds = folds % weight.shape[0]
+        selected_weights = weight[param_folds, units]  # (P, Ki)
+        unit_within = torch.distributions.Categorical(probs=selected_weights).sample()  # (P,)
+        return [ArityBranch(sample_ids, folds, unit_within) for _ in range(self.arity)]
+
 
 class TorchTensorDotLayer(TorchInnerLayer):
     r"""The tensor dot layer performs the following operations.

cirkit/backend/torch/queries.py

@@ -6,7 +6,14 @@
 from torch import Tensor
 
 from cirkit.backend.torch.circuits import TorchCircuit
-from cirkit.backend.torch.layers import TorchInnerLayer, TorchInputLayer, TorchLayer
+from cirkit.backend.torch.graph.modules import AddressBookEntry
+from cirkit.backend.torch.layers import (
+    ArityBranch,
+    BackwardSelection,
+    TorchInnerLayer,
+    TorchInputLayer,
+    TorchLayer,
+)
 from cirkit.utils.scope import Scope
 
 
@@ -187,14 +194,19 @@ def scopes_to_mask(
 class SamplingQuery(Query):
     """The sampling query object."""
 
-    def __init__(self, circuit: TorchCircuit) -> None:
+    def __init__(self, circuit: TorchCircuit, backward: bool = False) -> None:
         """Initialize a sampling query object. Currently, only sampling from the joint distribution
             is supported, i.e., sampling won't work in the case of circuits obtained by
             marginalization, or by observing evidence. Conditional sampling is currently not
             implemented.
 
         Args:
             circuit: The circuit to sample from.
+            backward: If True, use top-down (backward) ancestral sampling: walk the address
+                book in reverse from the root and at each layer track only the active paths
+                (via each layer's `backward_sample` method). If False (default), use the
+                bottom-up forward sampler that materializes samples for every (fold, unit)
+                of every layer.
 
         Raises:
             ValueError: If the circuit to sample from is not normalised.
@@ -207,6 +219,7 @@ def __init__(self, circuit: TorchCircuit) -> None:
         # TODO: add a check to verify the circuit is monotonic and normalized?
         super().__init__()
         self._circuit = circuit
+        self._backward = backward
 
     def __call__(self, num_samples: int = 1) -> tuple[Tensor, list[Tensor]]:
         """Sample a number of data points.
@@ -218,14 +231,18 @@ def __call__(self, num_samples: int = 1) -> tuple[Tensor, list[Tensor]]:
             A pair (samples, mixture_samples), consisting of (i) an assignment to the observed
             variables the circuit is defined on, and (ii) the samples of the finitely-discrete
             latent variables associated to the sum units. The samples (i) are returned as a
-            tensor of shape (num_samples, num_variables).
+            tensor of shape (num_samples, num_variables). In backward mode `mixture_samples`
+            is always an empty list.
 
         Raises:
             ValueError: if the number of samples is not a positive number.
         """
         if num_samples <= 0:
             raise ValueError("The number of samples must be a positive number")
 
+        if self._backward:
+            return _backward_sample(self._circuit, num_samples), []
+
         mixture_samples: list[Tensor] = []
         # samples: (O, K, num_samples, D)
         samples = self._circuit.evaluate(
@@ -273,3 +290,125 @@ def _pad_samples(self, samples: Tensor, scope_idx: Tensor) -> Tensor:
         fold_idx = torch.arange(samples.shape[0], device=samples.device)
         padded_samples[fold_idx, :, :, scope_idx.squeeze(dim=1)] = samples
         return padded_samples
+
+
+# --- Backward (top-down) sampling driver --------------------------------------------------
+#
+# Walks the circuit's address book in reverse. At each entry the per-layer `backward_sample`
+# method does the local sampling decision (sum/CPT: which input; input layer: write a value).
+# This driver handles cross-layer concerns: translating each arity-slot's `ArityBranch` into
+# the correct child entry via `entry.in_fold_idx[0]` and `entry.in_module_ids[0]`, including
+# the multi-child dispatch.
+
+
+@torch.no_grad()
+def _backward_sample(circuit: TorchCircuit, num_samples: int) -> Tensor:
+    device = next(circuit.parameters()).device
+    entries = list(circuit.address_book)
+
+    selections: dict[int, BackwardSelection] = {}
+
+    # Initialize root: all N samples start at fold=0, unit=0.
+    output_entry = entries[-1]
+    root_idx = output_entry.in_module_ids[0][0]
+    selections[root_idx] = (
+        torch.arange(num_samples, dtype=torch.long, device=device),
+        torch.zeros(num_samples, dtype=torch.long, device=device),
+        torch.zeros(num_samples, dtype=torch.long, device=device),
+    )
+
+    samples = torch.zeros(num_samples, circuit.num_variables, dtype=torch.long, device=device)
+
+    for entry_idx in range(len(entries) - 2, -1, -1):
+        entry = entries[entry_idx]
+        if entry.module is None or entry_idx not in selections:
+            continue
+        selection = selections[entry_idx]
+
+        if isinstance(entry.module, TorchInputLayer):
+            entry.module.backward_sample(selection, samples)
+        else:
+            assert isinstance(entry.module, TorchInnerLayer)
+            branches = entry.module.backward_sample(selection)
+            for h, branch in enumerate(branches):
+                _push_to_children(entry, branch, h, entries, selections)
+
+        del selections[entry_idx]
+
+    return samples
+
+
+def _push_to_children(
+    entry: AddressBookEntry[TorchLayer],
+    branch: ArityBranch,
+    h: int,
+    entries: list[AddressBookEntry],
+    selections: dict[int, BackwardSelection],
+) -> None:
+    """Translate `branch` (active at arity slot h) into a child entry selection."""
+    if branch.sample_ids.numel() == 0:
+        return
+
+    fold_idx_h = entry.in_fold_idx[0]
+    in_layer_ids = entry.in_module_ids[0]
+
+    # Compute the child concat-fold (index into the concatenation of all children's folds)
+    # for each path in `branch`, based on the shape of in_fold_idx.
+    if isinstance(fold_idx_h, tuple):
+        if fold_idx_h == (None,):
+            # unsqueeze dim=0: parent has 1 fold; arity slot h reads child concat-fold = h.
+            child_concat = torch.full_like(branch.folds, h)
+        else:
+            # (slice(None), None): arity is 1; concat-fold = parent fold.
+            child_concat = branch.folds
+    elif isinstance(fold_idx_h, Tensor):
+        if fold_idx_h.shape[1] == 1 and entry.module.arity == 1:
+            child_concat = fold_idx_h[branch.folds, 0]
+        else:
+            child_concat = fold_idx_h[branch.folds, h]
+    else:
+        raise RuntimeError(f"Unexpected in_fold_idx element type: {type(fold_idx_h)}")
+
+    # Dispatch the concat-fold across children.
+    if len(in_layer_ids) == 1:
+        _append_selection(
+            selections,
+            in_layer_ids[0],
+            branch.sample_ids,
+            child_concat,
+            branch.units,
+        )
+        return
+
+    offset = 0
+    for mid in in_layer_ids:
+        n_folds = entries[mid].module.num_folds
+        mask = (child_concat >= offset) & (child_concat < offset + n_folds)
+        if mask.any():
+            idx = mask.nonzero(as_tuple=True)[0]
+            _append_selection(
+                selections,
+                mid,
+                branch.sample_ids[idx],
+                child_concat[idx] - offset,
+                branch.units[idx],
+            )
+        offset += n_folds
+
+
+def _append_selection(
+    selections: dict[int, BackwardSelection],
+    mid: int,
+    sample_ids: Tensor,
+    folds: Tensor,
+    units: Tensor,
+) -> None:
+    if mid in selections:
+        old_s, old_f, old_u = selections[mid]
+        selections[mid] = (
+            torch.cat([old_s, sample_ids]),
+            torch.cat([old_f, folds]),
+            torch.cat([old_u, units]),
+        )
+    else:
+        selections[mid] = (sample_ids, folds, units)