Scan Morphism

ScanMorphism(cell) realizes the iterated Kleisli composition of a per-step cell over a sequence. ScanMorphism.rsample(x) applies the kernel at each step while threading the hidden state. ScanMorphism.log_joint(x, hidden_states) sums the per-step log density. It accepts the hidden-state trajectory either as a positional tensor or as a {state_key: tensor} dictionary. The default state_key is "h".

scan

Scan combinator: temporal recurrence over sequences.

A ScanMorphism wraps a recurrent cell and applies it across a sequence, threading hidden state from one time step to the next. This implements the standard RNN pattern:

h_t = cell(x_t, h_{t-1})

where cell : A * H -> H is a morphism (either a plain ContinuousMorphism or a MonadicProgram) whose domain is a product of the per-timestep input space A and the hidden state space H, and whose codomain is H.

Given a cell : A * H -> H, scan(cell) produces a morphism A -> H that, at runtime:

1. Expects a 3D input tensor of shape (batch, seq_len, dim_A).
2. Initializes hidden state h_0 (zeros or a learned parameter).
3. At each step t, concatenates x[:, t, :] with h to form the cell input, then calls cell.rsample to produce the new h.
4. Returns the final hidden state h_T of shape (batch, dim_H).

The scan's type in the categorical framework is:

scan(f : A x H -> H) : A -> H

where the sequence structure is implicit in the tensor's time dimension, following standard neural network conventions.

Initialization strategies

• "zeros": h_0 = 0 (default).
• "learned": h_0 is a learnable nn.Parameter.

Examples:

>>> from quivers.continuous.spaces import Euclidean, ProductSpace
>>> from quivers.continuous.families import ConditionalNormal
>>> A = Euclidean(name="input", dim=32)
>>> H = Euclidean(name="hidden", dim=64)
>>> cell = ConditionalNormal(ProductSpace(A, H), H, scale=0.1)
>>> scanned = ScanMorphism(cell, init="zeros")
>>> scanned.domain   # Euclidean(name="input", dim=32)
>>> scanned.codomain # Euclidean(name="hidden", dim=64)
>>> x = torch.randn(8, 10, 32)  # batch=8, seq_len=10, input_dim=32
>>> h = scanned.rsample(x)      # (8, 64)

ScanMorphism

ScanMorphism(cell: ContinuousMorphism, init: str = 'zeros')

Bases: ContinuousMorphism

Temporal scan: apply a recurrent cell across a sequence.

Wraps a cell morphism f : A * H -> H and produces a morphism A -> H that iterates over the time dimension of a 3D input tensor, threading hidden state forward.

This implements standard RNN-style recurrence::

h_0 = init
h_t = cell(concat(x_t, h_{t-1}))  for t = 1..T

The scan returns the final hidden state h_T.

PARAMETER	DESCRIPTION
cell	The recurrent cell. Must have a product domain A * H and codomain H, where H matches the last component of the product domain. TYPE: ContinuousMorphism
init	Initialization strategy for h_0. One of "zeros" (default) or "learned" (trainable initial state). TYPE: str DEFAULT: 'zeros'

Source code in src/quivers/continuous/scan.py

def __init__(self, cell: ContinuousMorphism, init: str = "zeros") -> None:
    input_space = _extract_input_space(cell)
    hidden_space = cell.codomain
    super().__init__(input_space, hidden_space)
    self._cell = cell
    self._init_strategy = init
    self._input_dim = _event_dim(input_space)
    self._hidden_dim = _event_dim(hidden_space)
    if init == "learned":
        self._h0 = nn.Parameter(torch.zeros(self._hidden_dim))
    elif init != "zeros":
        raise ValueError(
            f"unknown init strategy {init!r}; expected 'zeros' or 'learned'"
        )

rsample

rsample(x: Tensor, sample_shape: Size = Size()) -> Tensor

Run the cell across the time dimension of x.

PARAMETER	DESCRIPTION
x	Input sequence, in either layout _as_sequence reads: (batch, seq_len, input_dim) or the folded (batch, seq_len * input_dim). TYPE: Tensor
sample_shape	Additional leading sample dimensions (applied to the cell's rsample at the first time step only). TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Final hidden state. Shape (batch, hidden_dim), or (*sample_shape, batch, hidden_dim) if sample_shape is non-empty.

Source code in src/quivers/continuous/scan.py

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Run the cell across the time dimension of x.

    Parameters
    ----------
    x : torch.Tensor
        Input sequence, in either layout
        `_as_sequence`
        reads: ``(batch, seq_len, input_dim)`` or the folded
        ``(batch, seq_len * input_dim)``.
    sample_shape : torch.Size
        Additional leading sample dimensions (applied to the
        cell's rsample at the first time step only).

    Returns
    -------
    torch.Tensor
        Final hidden state. Shape ``(batch, hidden_dim)``,
        or ``(*sample_shape, batch, hidden_dim)`` if
        sample_shape is non-empty.
    """
    x = self._as_sequence(x)
    batch, seq_len, _ = x.shape
    h = self._initial_state(batch, x)
    for t in range(seq_len):
        x_t = x[:, t, :]
        cell_input = torch.cat([x_t, h], dim=-1)
        if t == 0 and len(sample_shape) > 0:
            h = self._cell.rsample(cell_input, sample_shape)
            h = self._flatten_cell_output(h)
            if len(sample_shape) > 0 and h.dim() > 2:
                x = x.unsqueeze(0).expand(*sample_shape, *x.shape)
        else:
            if h.dim() > 2:
                x_t = x[..., t, :]
                cell_input = torch.cat([x_t, h], dim=-1)
            h = self._cell.rsample(cell_input)
            h = self._flatten_cell_output(h)
    return h

base_dimension

base_dimension(x: Tensor) -> int | None

One cell's worth of coordinates per time step of the input.

The recurrence draws once per position, so its coordinate budget is the sequence length times the cell's own. The length is read off the input rather than declared, because scan(cell) : A -> H says nothing about how many positions a given input carries, and it is read through _as_sequence so a folded (batch, seq_len * input_dim) input reports the whole sequence's budget rather than one step's. Under-reporting it would hand push_base a block too short to run the recurrence on, and a chain that sliced its coordinates by that count would give every later factor the wrong ones.

Source code in src/quivers/continuous/scan.py

def base_dimension(self, x: torch.Tensor) -> int | None:
    """One cell's worth of coordinates per time step of the input.

    The recurrence draws once per position, so its coordinate
    budget is the sequence length times the cell's own. The length
    is read off the input rather than declared, because
    ``scan(cell) : A -> H`` says nothing about how many positions a
    given input carries, and it is read through
    `_as_sequence`
    so a folded ``(batch, seq_len * input_dim)`` input reports the
    whole sequence's budget rather than one step's. Under-reporting
    it would hand
    [`push_base`][quivers.continuous.scan.ScanMorphism.push_base]
    a block too short to run the recurrence on, and a chain that
    sliced its coordinates by that count would give every later
    factor the wrong ones.
    """
    seq = self._as_sequence(x)
    step = self._step_dimension(seq)
    if step is None:
        return None
    return int(seq.shape[1]) * step

push_base

push_base(x: Tensor, base: Tensor) -> Tensor

Run the recurrence on supplied coordinates instead of draws.

Time step t reads the t-th block of base, so the trajectory is a deterministic function of the coordinates and the input, and the same coordinates always produce the same final state.

Source code in src/quivers/continuous/scan.py

def push_base(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
    """Run the recurrence on supplied coordinates instead of draws.

    Time step ``t`` reads the ``t``-th block of ``base``, so the
    trajectory is a deterministic function of the coordinates and
    the input, and the same coordinates always produce the same
    final state.
    """
    seq = self._as_sequence(x)
    step = self._step_dimension(seq)
    if step is None:
        raise ValueError(
            f"ScanMorphism.push_base: the cell "
            f"{type(self._cell).__name__} declares no "
            f"reparameterization, so the recurrence has none either."
        )
    batch, seq_len, _ = seq.shape
    h = self._initial_state(batch, seq)
    for t in range(seq_len):
        cell_input = torch.cat([seq[:, t, :], h], dim=-1)
        block = base[:, t * step : (t + 1) * step]
        h = self._flatten_cell_output(self._cell.push_base(cell_input, block))
    return h

reference_trajectory

reference_trajectory(x: Tensor, y: Tensor) -> Tensor

The states :math:h_1, \ldots, h_T this kernel scores at y.

The recurrence's own deterministic skeleton, re-anchored at the observed final state: for :math:t < T the state is the cell's image of the base measure's origin, and :math:h_T is y.

Prefix states use push_base at zero coordinates; the final state is y. The resulting trajectory is deterministic.

PARAMETER	DESCRIPTION
x	Input sequence, in either layout _as_sequence reads. TYPE: Tensor
y	Observed final state. Shape (batch, hidden_dim). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Shape (batch, seq_len, hidden_dim).

RAISES	DESCRIPTION
ValueError	If the cell declares no reparameterization, so the prefix has no canonical states to take.

Source code in src/quivers/continuous/scan.py

def reference_trajectory(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """The states :math:`h_1, \\ldots, h_T` this kernel scores at ``y``.

    The recurrence's own deterministic skeleton, re-anchored at the
    observed final state: for :math:`t < T` the state is the cell's
    image of the base measure's origin, and :math:`h_T` is ``y``.

    Prefix states use ``push_base`` at zero coordinates; the final
    state is ``y``. The resulting trajectory is deterministic.

    Parameters
    ----------
    x : torch.Tensor
        Input sequence, in either layout
        `_as_sequence`
        reads.
    y : torch.Tensor
        Observed final state. Shape ``(batch, hidden_dim)``.

    Returns
    -------
    torch.Tensor
        Shape ``(batch, seq_len, hidden_dim)``.

    Raises
    ------
    ValueError
        If the cell declares no reparameterization, so the prefix
        has no canonical states to take.
    """
    seq = self._as_sequence(x)
    batch, seq_len, _ = seq.shape
    h = self._initial_state(batch, seq)
    states: list[torch.Tensor] = []
    for t in range(seq_len - 1):
        cell_input = torch.cat([seq[:, t, :], h], dim=-1)
        dimension = self._cell.base_dimension(cell_input)
        if dimension is None:
            raise ValueError(
                f"ScanMorphism.reference_trajectory: the cell "
                f"{type(self._cell).__name__} declares no "
                f"reparameterization, so its base measure has no "
                f"origin and the recurrence has no canonical "
                f"prefix to score along."
            )
        base = torch.zeros(batch, dimension, device=seq.device, dtype=h.dtype)
        h = self._flatten_cell_output(self._cell.push_base(cell_input, base))
        states.append(h)
    states.append(y.reshape(batch, self._hidden_dim))
    return torch.stack(states, dim=1)

log_prob

log_prob(x: Tensor, y: Tensor) -> Tensor

Log-density of the scan's trajectory ending at y.

scan(cell) denotes a Kleisli morphism :math:\mathbf{x}_{1:T} \to \mathcal{G}(h_T) whose density at :math:h_T marginalizes every intermediate state:

.. math::

p(h_T \mid x_{1:T}) = \int
p(h_T \mid x_T, h_{T-1})
\prod_{t<T} p(h_t \mid x_t, h_{t-1})
\, dh_{1:T-1}.

The implementation scores the fixed trajectory returned by reference_trajectory with log_joint:

.. math::

\sum_{t=1}^{T} \log p(h_t \mid x_t, h_{t-1}),
\qquad h_T = y .

Floating-point reassociation in the recurrent prefix may be amplified by later transitions. Numerical comparisons should use the same cell implementation.

This is the joint density of the fixed trajectory, not the marginal density of its endpoint.

If the cell has no conditional density, this method returns a zero contribution.

PARAMETER	DESCRIPTION
x	Input sequence, in either layout _as_sequence reads. TYPE: Tensor
y	Final hidden state. Shape (batch, hidden_dim). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Shape (batch,).

Source code in src/quivers/continuous/scan.py

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-density of the scan's trajectory ending at ``y``.

    ``scan(cell)`` denotes a Kleisli morphism
    :math:`\\mathbf{x}_{1:T} \\to \\mathcal{G}(h_T)` whose density at
    :math:`h_T` marginalizes every intermediate state:

    .. math::

        p(h_T \\mid x_{1:T}) = \\int
        p(h_T \\mid x_T, h_{T-1})
        \\prod_{t<T} p(h_t \\mid x_t, h_{t-1})
        \\, dh_{1:T-1}.

    The implementation scores the fixed trajectory returned by
    ``reference_trajectory`` with ``log_joint``:

    .. math::

        \\sum_{t=1}^{T} \\log p(h_t \\mid x_t, h_{t-1}),
        \\qquad h_T = y .

    Floating-point reassociation in the recurrent prefix may be
    amplified by later transitions. Numerical comparisons should
    use the same cell implementation.

    This is the joint density of the fixed trajectory, not the
    marginal density of its endpoint.

    If the cell has no conditional density, this method returns a
    zero contribution.

    Parameters
    ----------
    x : torch.Tensor
        Input sequence, in either layout
        `_as_sequence`
        reads.
    y : torch.Tensor
        Final hidden state. Shape ``(batch, hidden_dim)``.

    Returns
    -------
    torch.Tensor
        Shape ``(batch,)``.
    """
    seq = self._as_sequence(x)
    if not self._cell.has_conditional_density():
        dtype = seq.dtype if seq.is_floating_point() else torch.get_default_dtype()
        return torch.zeros(seq.shape[0], device=seq.device, dtype=dtype)
    return self.log_joint(seq, self.reference_trajectory(seq, y))

log_joint

log_joint(x: Tensor, hidden_states: 'torch.Tensor | dict[str, torch.Tensor]', *, state_key: str = 'h') -> Tensor

Joint log-density given all intermediate hidden states.

Computes: log p(h_1, ..., h_T | x_{1:T}) = sum_t log p(h_t | x_t, h_{t-1})

PARAMETER	DESCRIPTION
x	Input sequence, in either layout _as_sequence reads. TYPE: Tensor
hidden_states	All hidden states including final, shape (batch, seq_len, hidden_dim). May be passed positionally as a tensor or via a dict keyed by state_key (so the inference layer's standard log_joint(x, observations: dict) contract works without an adapter). TYPE: Tensor | dict[str, Tensor]
state_key	Dict key under which the hidden-state tensor is looked up when hidden_states is a dict. Defaults to "h". TYPE: str DEFAULT: 'h'

RETURNS	DESCRIPTION
Tensor	Joint log-density. Shape (batch,).

Source code in src/quivers/continuous/scan.py

def log_joint(
    self,
    x: torch.Tensor,
    hidden_states: "torch.Tensor | dict[str, torch.Tensor]",
    *,
    state_key: str = "h",
) -> torch.Tensor:
    """Joint log-density given all intermediate hidden states.

    Computes:
        log p(h_1, ..., h_T | x_{1:T}) =
            sum_t log p(h_t | x_t, h_{t-1})

    Parameters
    ----------
    x : torch.Tensor
        Input sequence, in either layout
        `_as_sequence`
        reads.
    hidden_states : torch.Tensor | dict[str, torch.Tensor]
        All hidden states including final, shape
        ``(batch, seq_len, hidden_dim)``. May be passed
        positionally as a tensor or via a dict keyed by
        ``state_key`` (so the inference layer's standard
        ``log_joint(x, observations: dict)`` contract works
        without an adapter).
    state_key : str
        Dict key under which the hidden-state tensor is
        looked up when ``hidden_states`` is a dict. Defaults
        to ``"h"``.

    Returns
    -------
    torch.Tensor
        Joint log-density. Shape ``(batch,)``.
    """
    if isinstance(hidden_states, dict):
        hidden_states = hidden_states[state_key]
    seq = self._as_sequence(x)
    batch, seq_len, _ = seq.shape
    states = hidden_states.reshape(batch, seq_len, self._hidden_dim)
    h = self._initial_state(batch, seq)
    total = torch.zeros(batch, device=seq.device, dtype=h.dtype)
    for t in range(seq_len):
        h_t = states[:, t, :]
        cell_input = torch.cat([seq[:, t, :], h], dim=-1)
        total = total + self._cell.log_prob(cell_input, h_t)
        h = h_t
    return total