Continuous Morphisms

Operational kernels over discrete or continuous carriers. A ContinuousMorphism evaluates log densities and draws samples rather than materializing a finite morphism tensor.

morphisms

Markov kernels on discrete, continuous, and mixed spaces.

ContinuousMorphism defines log_prob and rsample. The >> operator composes kernels and @ forms their independent product. Discrete intermediates are marginalized by finite summation. Continuous intermediates are scored along the deterministic reference path described by SampledComposition.log_prob.

ContinuousMorphism

ContinuousMorphism(domain: AnySpace, codomain: AnySpace)

Bases: Module, ABC

Abstract base for morphisms involving continuous spaces.

Subclasses must implement log_prob and rsample. The composition operator >> and product operator @ are provided and dispatch to SampledComposition and ProductContinuousMorphism respectively.

Unlike discrete Morphism (which materializes a full tensor), ContinuousMorphism is defined operationally: it can evaluate log-densities and generate reparameterized samples.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. TYPE: SetObject or ContinuousSpace

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

def __init__(self, domain: AnySpace, codomain: AnySpace) -> None:
    super().__init__()
    self._domain = domain
    self._codomain = codomain

domain property

domain: AnySpace

Source space.

codomain property

codomain: AnySpace

Target space.

support property

support: Constraint

The support constraint of the distribution this morphism samples from, in the form of a torch.distributions.constraints.Constraint.

Used by variational guides (quivers.inference.AutoNormalGuide, quivers.inference.AutoDeltaGuide) to determine the correct bijector that maps an unconstrained variational approximation back to the constrained support of the prior, so that samples used to evaluate the prior's log_prob lie inside its support (avoiding Expected value to be within the support of the distribution errors).

Subclasses representing a constrained distribution family (HalfNormal, Beta, Uniform, Dirichlet, LogitNormal, Wishart, …) should override this property to return the appropriate constraint. The default is torch.distributions.constraints.real, which is correct for unconstrained families like Normal and discrete codomains (where the guide skips the site anyway).

log_prob abstractmethod

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

Log-probability (density) of y given x.

PARAMETER	DESCRIPTION
x	Inputs. Shape (batch,) for discrete domain or (batch, domain_dim) for continuous domain. TYPE: Tensor
y	Outputs. Shape (batch,) for discrete codomain or (batch, codomain_dim) for continuous codomain. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Log-probabilities/densities. Shape (batch,).

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

@abstractmethod
def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability (density) of y given x.

    Parameters
    ----------
    x : torch.Tensor
        Inputs. Shape (batch,) for discrete domain or
        (batch, domain_dim) for continuous domain.
    y : torch.Tensor
        Outputs. Shape (batch,) for discrete codomain or
        (batch, codomain_dim) for continuous codomain.

    Returns
    -------
    torch.Tensor
        Log-probabilities/densities. Shape (batch,).
    """
    ...

rsample abstractmethod

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

Reparameterized samples from p(. | x).

Gradients flow through the returned samples back to the parameters of this morphism (and to x if the domain is continuous).

PARAMETER	DESCRIPTION
x	Inputs. Shape (batch,) or (batch, domain_dim). TYPE: Tensor
sample_shape	Additional leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Samples. Shape (sample_shape, batch, codomain_dim) for continuous codomain, or (sample_shape, batch) for discrete.

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

@abstractmethod
def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Reparameterized samples from p(. | x).

    Gradients flow through the returned samples back to the
    parameters of this morphism (and to x if the domain is
    continuous).

    Parameters
    ----------
    x : torch.Tensor
        Inputs. Shape (batch,) or (batch, domain_dim).
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples. Shape (*sample_shape, batch, codomain_dim) for
        continuous codomain, or (*sample_shape, batch) for discrete.
    """
    ...

has_conditional_density

has_conditional_density() -> bool

Whether log_prob evaluates a density rather than raising.

Every kernel whose conditional law is a named family answers yes. A kernel that denotes a program answers no: its density at a value marginalizes the program's internal draws, which no closed form covers, and its log_prob says so by raising.

A caller deciding between two constructions needs that answer before it calls, not as an exception afterwards, so the capability is declared rather than discovered. The parallel with point_mass_value and base_dimension is exact: each reports a structural property of the kernel that determines which exact treatment is open to a caller, and none of them is a probe by trial.

RETURNS	DESCRIPTION
bool	True for a kernel with an evaluable conditional density.

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

def has_conditional_density(self) -> bool:
    """Whether
    [`log_prob`][quivers.continuous.morphisms.ContinuousMorphism.log_prob]
    evaluates a density rather than raising.

    Every kernel whose conditional law is a named family answers
    yes. A kernel that denotes a *program* answers no: its density
    at a value marginalizes the program's internal draws, which no
    closed form covers, and its `log_prob` says so by raising.

    A caller deciding between two constructions needs that answer
    before it calls, not as an exception afterwards, so the
    capability is declared rather than discovered. The parallel
    with
    [`point_mass_value`][quivers.continuous.morphisms.ContinuousMorphism.point_mass_value]
    and
    [`base_dimension`][quivers.continuous.morphisms.ContinuousMorphism.base_dimension]
    is exact: each reports a structural property of the kernel that
    determines which exact treatment is open to a caller, and none
    of them is a probe by trial.

    Returns
    -------
    bool
        True for a kernel with an evaluable conditional density.
    """
    return True

point_mass_value

point_mass_value(x: Tensor) -> Tensor | None

The single value this kernel puts all of its mass on, or None.

A morphism whose conditional law is a Dirac delta :math:\delta_{T(x)} returns :math:T(x); every other morphism returns None. The distinction is what lets SampledComposition collapse an integral over a degenerate intermediate to a single evaluation, which is exact rather than approximate.

PARAMETER	DESCRIPTION
x	Inputs. Shape (batch,) or (batch, domain_dim). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor or None	The deterministic image of x, or None when the kernel is genuinely stochastic.

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

def point_mass_value(self, x: torch.Tensor) -> torch.Tensor | None:
    """The single value this kernel puts all of its mass on, or None.

    A morphism whose conditional law is a Dirac delta
    :math:`\\delta_{T(x)}` returns :math:`T(x)`; every other
    morphism returns ``None``. The distinction is what lets
    [`SampledComposition`][quivers.continuous.morphisms.SampledComposition]
    collapse an integral over a degenerate intermediate to a
    single evaluation, which is exact rather than approximate.

    Parameters
    ----------
    x : torch.Tensor
        Inputs. Shape ``(batch,)`` or ``(batch, domain_dim)``.

    Returns
    -------
    torch.Tensor or None
        The deterministic image of ``x``, or ``None`` when the
        kernel is genuinely stochastic.
    """
    del x
    return None

base_dimension

base_dimension(x: Tensor) -> int | None

Standard-normal coordinates this kernel's reparameterization reads.

A morphism that can be written :math:y = T_x(\varepsilon) with :math:\varepsilon standard normal reports how many coordinates :math:T_x consumes at this input; every other morphism reports None. A deterministic map consumes none and reports 0.

The count may depend on x: an embedding kernel reading a (batch, seq) index matrix places one Gaussian per position, so it consumes seq * dim coordinates where the same kernel on a (batch,) index vector consumes dim.

PARAMETER	DESCRIPTION
x	Conditioning inputs. Shape (batch,) or (batch, domain_dim). TYPE: Tensor

RETURNS	DESCRIPTION
int or None	The coordinate count, or None when this morphism has no reparameterization to offer.

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

def base_dimension(self, x: torch.Tensor) -> int | None:
    """Standard-normal coordinates this kernel's reparameterization reads.

    A morphism that can be written :math:`y = T_x(\\varepsilon)`
    with :math:`\\varepsilon` standard normal reports how many
    coordinates :math:`T_x` consumes at this input; every other
    morphism reports ``None``. A deterministic map consumes none
    and reports ``0``.

    The count may depend on ``x``: an embedding kernel reading a
    ``(batch, seq)`` index matrix places one Gaussian per position,
    so it consumes ``seq * dim`` coordinates where the same kernel
    on a ``(batch,)`` index vector consumes ``dim``.

    Parameters
    ----------
    x : torch.Tensor
        Conditioning inputs. Shape ``(batch,)`` or
        ``(batch, domain_dim)``.

    Returns
    -------
    int or None
        The coordinate count, or ``None`` when this morphism has
        no reparameterization to offer.
    """
    del x
    if type(self).point_mass_value is not ContinuousMorphism.point_mass_value:
        return 0
    return None

push_base

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

Push standard-normal coordinates through the reparameterization.

Evaluates :math:T_x(\varepsilon) for the map base_dimension describes. The map is a pure function of (x, base): it reads no random state, which is what lets a caller build a quadrature out of it and get the same nodes on every call.

The default covers the degenerate case, where the map ignores its (empty) coordinates and returns the point mass.

PARAMETER	DESCRIPTION
x	Conditioning inputs. Shape (batch, *domain). TYPE: Tensor
base	Standard-normal coordinates. Shape (batch, dimension) for the dimension base_dimension reports at x. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Shape (batch, *event).

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

def push_base(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
    """Push standard-normal coordinates through the reparameterization.

    Evaluates :math:`T_x(\\varepsilon)` for the map
    [`base_dimension`][quivers.continuous.morphisms.ContinuousMorphism.base_dimension]
    describes. The map is a pure function of ``(x, base)``: it
    reads no random state, which is what lets a caller build a
    quadrature out of it and get the same nodes on every call.

    The default covers the degenerate case, where the map ignores
    its (empty) coordinates and returns the point mass.

    Parameters
    ----------
    x : torch.Tensor
        Conditioning inputs. Shape ``(batch, *domain)``.
    base : torch.Tensor
        Standard-normal coordinates. Shape ``(batch, dimension)``
        for the dimension `base_dimension` reports at ``x``.

    Returns
    -------
    torch.Tensor
        Shape ``(batch, *event)``.
    """
    del base
    value = self.point_mass_value(x)
    if value is None:
        raise ValueError(
            f"{type(self).__name__}.push_base: this morphism "
            f"declares no reparameterization, so there is nothing "
            f"to push coordinates through. Override "
            f"`base_dimension` and `push_base` together, or leave "
            f"both at their defaults so callers see the absence "
            f"rather than a wrong value."
        )
    return value

marginal_quadrature

marginal_quadrature(x: Tensor, count: int) -> tuple[Tensor, Tensor] | None

A deterministic rule for integrating against p(. | x).

Returns (nodes, log_weights) approximating

.. math::

\int p(y \mid x) \, \varphi(y) \, dy
\;\approx\;
\sum_i \exp(\log w_i) \, \varphi(y_i)

Point masses return one unit-weight node. Reparameterized kernels return equally weighted Sobol nodes produced by push_base. Kernels without either representation return None. This method does not consume random state.

PARAMETER	DESCRIPTION
x	Conditioning inputs. Shape (batch,) or (batch, domain_dim). TYPE: Tensor
count	Requested number of nodes. An implementation may return fewer (an exact rule needs one) or round up to the count its construction is balanced at. TYPE: int

RETURNS	DESCRIPTION
tuple[Tensor, Tensor] or None	Nodes and log-weights, or None when this morphism provides no deterministic rule.

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

def marginal_quadrature(
    self, x: torch.Tensor, count: int
) -> tuple[torch.Tensor, torch.Tensor] | None:
    """A deterministic rule for integrating against ``p(. | x)``.

    Returns ``(nodes, log_weights)`` approximating

    .. math::

        \\int p(y \\mid x) \\, \\varphi(y) \\, dy
        \\;\\approx\\;
        \\sum_i \\exp(\\log w_i) \\, \\varphi(y_i)

    Point masses return one unit-weight node. Reparameterized
    kernels return equally weighted Sobol nodes produced by
    ``push_base``. Kernels without either representation return
    ``None``. This method does not consume random state.

    Parameters
    ----------
    x : torch.Tensor
        Conditioning inputs. Shape ``(batch,)`` or
        ``(batch, domain_dim)``.
    count : int
        Requested number of nodes. An implementation may return
        fewer (an exact rule needs one) or round up to the count
        its construction is balanced at.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor] or None
        Nodes and log-weights, or ``None`` when this morphism
        provides no deterministic rule.
    """
    value = self.point_mass_value(x)
    if value is not None:
        nodes = value.unsqueeze(0)
        log_weights = torch.zeros(1, device=nodes.device, dtype=nodes.dtype)
        return nodes, log_weights
    dimension = self.base_dimension(x)
    if dimension is None:
        return None
    batch = x.shape[0]
    dtype = x.dtype if x.is_floating_point() else torch.get_default_dtype()
    base = sobol_normal_points(dimension, count, x.device, dtype)
    n = base.shape[0]
    x_rows = x.unsqueeze(0).expand(n, *x.shape).reshape(n * batch, *x.shape[1:])
    base_rows = (
        base.unsqueeze(1).expand(n, batch, dimension).reshape(n * batch, dimension)
    )
    pushed = self.push_base(x_rows, base_rows)
    nodes = pushed.reshape(n, batch, *pushed.shape[1:])
    log_weights = torch.full(
        (n,), -math.log(float(n)), device=nodes.device, dtype=nodes.dtype
    )
    return nodes, log_weights

sample

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

Non-reparameterized samples (no gradient through samples).

PARAMETER	DESCRIPTION
x	Inputs. TYPE: Tensor
sample_shape	Additional leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Samples (detached from computation graph).

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

def sample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Non-reparameterized samples (no gradient through samples).

    Parameters
    ----------
    x : torch.Tensor
        Inputs.
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples (detached from computation graph).
    """
    with torch.no_grad():
        return self.rsample(x, sample_shape)

__rshift__

__rshift__(other: object) -> ContinuousMorphism

Composition via ancestral sampling: self >> other.

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

def __rshift__(self, other: object) -> ContinuousMorphism:
    """Composition via ancestral sampling: self >> other."""
    if isinstance(other, ContinuousMorphism):
        return SampledComposition(self, other)
    from quivers.core.morphisms import Morphism

    if isinstance(other, Morphism):
        return SampledComposition(self, DiscreteAsContinuous(other))
    return NotImplemented

__rrshift__

__rrshift__(other: object) -> ContinuousMorphism

Handle discrete_morphism >> continuous_morphism.

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

def __rrshift__(self, other: object) -> ContinuousMorphism:
    """Handle discrete_morphism >> continuous_morphism."""
    from quivers.core.morphisms import Morphism

    if isinstance(other, Morphism):
        return SampledComposition(DiscreteAsContinuous(other), self)
    return NotImplemented

__matmul__

__matmul__(other: object) -> ProductContinuousMorphism

Independent product: self @ other.

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

def __matmul__(self, other: object) -> ProductContinuousMorphism:
    """Independent product: self @ other."""
    if isinstance(other, ContinuousMorphism):
        return ProductContinuousMorphism(self, other)
    from quivers.core.morphisms import Morphism

    if isinstance(other, Morphism):
        return ProductContinuousMorphism(self, DiscreteAsContinuous(other))
    return NotImplemented

MarginalizedFactor

MarginalizedFactor(base: ContinuousMorphism)

Bases: ContinuousMorphism

Score-suppressed wrapper for a marginalized block's live sites.

An ungrouped marginalize block keeps its latent draw and its terminal observe as live sites so a forward trace still produces the sampled coordinate and response (ancestral sampling and synthetic-data generation both read those sites). Their densities, however, are carried once by the block's integrated score step, so adding them to the joint again would double-count the very factors the marginal already integrates. This wrapper delegates sampling to the base morphism yet reports a zero log-density, keeping the joint free of the raw per-draw factor while preserving forward behaviour.

PARAMETER	DESCRIPTION
base	The underlying family whose sampling behaviour is preserved. TYPE: ContinuousMorphism

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

def __init__(self, base: ContinuousMorphism) -> None:
    super().__init__(base.domain, base.codomain)
    self.base = base

support property

support: Constraint

Delegate the support constraint to the wrapped family.

rsample

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

Sample from the base family (forward behaviour is preserved).

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample from the base family (forward behaviour is preserved)."""
    return self.base.rsample(x, sample_shape)

log_prob

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

Report a zero log-density.

The factor's density is carried by the block's integrated score step; contributing it here would double-count it in the joint.

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Report a zero log-density.

    The factor's density is carried by the block's integrated score
    step; contributing it here would double-count it in the joint.
    """
    del x
    return torch.zeros((), device=y.device, dtype=torch.get_default_dtype())

SampledComposition

SampledComposition(left: ContinuousMorphism, right: ContinuousMorphism, n_intermediate: int = 100)

Bases: ContinuousMorphism

Composition of morphisms via ancestral sampling.

rsample draws from left and then right. log_prob sums over a discrete intermediate. For a stochastic continuous intermediate it scores the deterministic reference path, not the endpoint marginal.

PARAMETER	DESCRIPTION
left	First morphism (applied first). TYPE: ContinuousMorphism
right	Second morphism (applied second). TYPE: ContinuousMorphism
n_intermediate	Node count for the deterministic rule this composition offers through marginal_quadrature to a caller that asks for one. The composite density does not ask: it scores the canonical path and integrates nothing, so this count does not reach log_prob. TYPE: int DEFAULT: 100

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

def __init__(
    self,
    left: ContinuousMorphism,
    right: ContinuousMorphism,
    n_intermediate: int = 100,
) -> None:
    super().__init__(left.domain, right.codomain)
    self.left = left
    self.right = right
    self.n_intermediate = n_intermediate

factors property

factors: tuple[ContinuousMorphism, ...]

The composition flattened into its non-composite factors.

(a >> b) >> c and a >> (b >> c) both report (a, b, c): association is invisible to the kernel the composition denotes, and every intermediate between adjacent factors is an object the chain integrates over. A caller that wants those intermediates as named sites walks this tuple.

base_dimension

base_dimension(x: Tensor) -> int | None

Total coordinates the whole chain's reparameterization reads.

A chain is reparameterized by reparameterizing each factor and threading the result forward, so its coordinate budget is the sum of its factors'. One factor without a reparameterization leaves the chain without one.

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

def base_dimension(self, x: torch.Tensor) -> int | None:
    """Total coordinates the whole chain's reparameterization reads.

    A chain is reparameterized by reparameterizing each factor and
    threading the result forward, so its coordinate budget is the
    sum of its factors'. One factor without a reparameterization
    leaves the chain without one.
    """
    dimensions = chain_dimensions(self.factors, x)
    if dimensions is None:
        return None
    return sum(dimensions)

push_base

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

Thread the coordinates through the chain, factor by factor.

Each factor consumes its own contiguous block of base, so no two factors share a coordinate.

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

def push_base(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
    """Thread the coordinates through the chain, factor by factor.

    Each factor consumes its own contiguous block of ``base``, so
    no two factors share a coordinate.
    """
    dimensions = chain_dimensions(self.factors, x)
    if dimensions is None:
        raise ValueError(
            f"SampledComposition.push_base: a factor of this chain "
            f"declares no reparameterization, so the chain has "
            f"none either. The factors are "
            f"{[type(f).__name__ for f in self.factors]!r}."
        )
    return chain_push_base(self.factors, x, base, dimensions)

rsample

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

Ancestral sampling: y ~ f(x, .), then z ~ g(y, .).

PARAMETER	DESCRIPTION
x	Inputs to the composition. TYPE: Tensor
sample_shape	Additional sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Samples from the composed morphism.

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Ancestral sampling: y ~ f(x, .), then z ~ g(y, .).

    Parameters
    ----------
    x : torch.Tensor
        Inputs to the composition.
    sample_shape : torch.Size
        Additional sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples from the composed morphism.
    """
    y = self.left.rsample(x, sample_shape)
    if len(sample_shape) > 0:
        leading = y.shape[: len(sample_shape)]
        batch = x.shape[0]
        flat_size = int(torch.tensor(leading).prod().item()) * batch
        if y.dim() > len(sample_shape) + 1:
            event_dims = y.shape[len(sample_shape) + 1 :]
            flat_y = y.reshape(flat_size, *event_dims)
        else:
            flat_y = y.reshape(flat_size)
    else:
        flat_y = y
    z = self.right.rsample(flat_y)
    if len(sample_shape) > 0:
        batch = x.shape[0]
        if z.dim() > 1:
            event_dims = z.shape[1:]
            z = z.reshape(*sample_shape, batch, *event_dims)
        else:
            z = z.reshape(*sample_shape, batch)
    return z

log_prob

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

Log-probability of y given x through the composition.

A discrete intermediate is marginalized exactly, by finite summation over its elements. A continuous one is not marginalized at all: the chain is scored along the canonical path _log_prob_reference_path describes, which is exact where an integral would have been approximate, and is a pure function of (x, y) where a rule would have made it a function of the node count as well.

Both branches return the same number when every intermediate is degenerate, which is the case the two readings share.

PARAMETER	DESCRIPTION
x	Inputs. Shape (batch,) or (batch, dom_dim). TYPE: Tensor
y	Outputs. Shape (batch,) or (batch, cod_dim). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Log-probabilities. Shape (batch,).

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability of y given x through the composition.

    A discrete intermediate is marginalized exactly, by finite
    summation over its elements. A continuous one is not
    marginalized at all: the chain is scored along the canonical
    path
    `_log_prob_reference_path`
    describes, which is exact where an integral would have been
    approximate, and is a pure function of ``(x, y)`` where a rule
    would have made it a function of the node count as well.

    Both branches return the same number when every intermediate
    is degenerate, which is the case the two readings share.

    Parameters
    ----------
    x : torch.Tensor
        Inputs. Shape (batch,) or (batch, dom_dim).
    y : torch.Tensor
        Outputs. Shape (batch,) or (batch, cod_dim).

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape (batch,).
    """
    intermediate = self.left.codomain
    if isinstance(intermediate, SetObject):
        return self._log_prob_exact(x, y, intermediate)
    else:
        return self._log_prob_reference_path(x, y)

ProductContinuousMorphism

ProductContinuousMorphism(left: ContinuousMorphism, right: ContinuousMorphism)

Bases: ContinuousMorphism

Independent product of two continuous morphisms.

Given f: A -> B and g: C -> D, produces f @ g: (A, C) -> (B, D) where p_{f@g}((y,z) | (x,w)) = f(y | x) * g(z | w).

Domain inputs are concatenated: (x, w) as a single vector. Codomain outputs are concatenated: (y, z) as a single vector. For discrete components, indices are embedded as 1-d floats.

PARAMETER	DESCRIPTION
left	Left factor morphism. TYPE: ContinuousMorphism
right	Right factor morphism. TYPE: ContinuousMorphism

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

def __init__(self, left: ContinuousMorphism, right: ContinuousMorphism) -> None:
    dom = _combine_spaces(left.domain, right.domain)
    cod = _combine_spaces(left.codomain, right.codomain)
    super().__init__(dom, cod)
    self.left = left
    self.right = right
    self._left_dom_dim = _event_dim(left.domain)
    self._right_dom_dim = _event_dim(right.domain)
    self._left_cod_dim = _event_dim(left.codomain)
    self._right_cod_dim = _event_dim(right.codomain)

base_dimension

base_dimension(x: Tensor) -> int | None

Sum of the two factors' coordinate budgets at their own inputs.

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

def base_dimension(self, x: torch.Tensor) -> int | None:
    """Sum of the two factors' coordinate budgets at their own inputs."""
    x_left, x_right = self._split_input(dimension_probe(x))
    left = self.left.base_dimension(x_left)
    right = self.right.base_dimension(x_right)
    if left is None or right is None:
        return None
    return left + right

push_base

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

Push each factor's own coordinate block through that factor.

The factors are independent given the input, so the product's reparameterization is the pair of theirs on disjoint coordinate blocks, concatenated along the feature axis exactly as rsample concatenates its draws.

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

def push_base(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
    """Push each factor's own coordinate block through that factor.

    The factors are independent given the input, so the product's
    reparameterization is the pair of theirs on disjoint
    coordinate blocks, concatenated along the feature axis exactly
    as `rsample` concatenates its draws.
    """
    x_left, x_right = self._split_input(x)
    probe_left, probe_right = self._split_input(dimension_probe(x))
    left_dimension = self.left.base_dimension(probe_left)
    right_dimension = self.right.base_dimension(probe_right)
    if left_dimension is None or right_dimension is None:
        raise ValueError(
            f"ProductContinuousMorphism.push_base: factor "
            f"{type(self.left).__name__} @ "
            f"{type(self.right).__name__} declares no "
            f"reparameterization, so the product has none either."
        )
    y_left = self.left.push_base(x_left, base[:, :left_dimension])
    y_right = self.right.push_base(x_right, base[:, left_dimension:])
    if y_left.dim() < y_right.dim():
        y_left = y_left.unsqueeze(-1)
    elif y_right.dim() < y_left.dim():
        y_right = y_right.unsqueeze(-1)
    return torch.cat([y_left, y_right], dim=-1)

FanOutMorphism

FanOutMorphism(components: list)

Bases: ContinuousMorphism

Fan-out morphism: copy input to N morphisms, concatenate outputs.

Given f_1: A -> B_1, f_2: A -> B_2, ..., f_N: A -> B_N, produces fan(f_1, ..., f_N): A -> B_1 * B_2 * ... * B_N where the input A is copied to all N morphisms.

Unlike the tensor product (f @ g), which takes a product domain (A * C), fan-out feeds the same input to all morphisms. This implements the diagonal morphism Delta: A -> A^N followed by the product f_1 @ f_2 @ ... @ f_N.

PARAMETER	DESCRIPTION
components	The morphisms to fan out to. All must share the same domain. TYPE: list[ContinuousMorphism]

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

def __init__(self, components: list) -> None:
    from quivers.core.morphisms import Morphism as _CatMorphism

    if not components:
        raise ValueError("fan-out requires at least one component")
    # Backend-agnostic V-Cat morphisms (those that aren't
    # already ContinuousMorphism subclasses) get wrapped in a
    # deterministic continuous adapter so the FanOut's rsample
    # / log_prob loop can dispatch uniformly. The wrapping
    # exposes the V-Cat tensor through a categorical
    # ``rsample`` that gathers / contracts the tensor against
    # the input; ``log_prob`` evaluates the V-Cat tensor as a
    # categorical likelihood when meaningful.
    wrapped_components: list[ContinuousMorphism] = []
    for c in components:
        if isinstance(c, ContinuousMorphism):
            wrapped_components.append(c)
        elif isinstance(c, _CatMorphism):
            wrapped_components.append(DiscreteAsContinuous(c))
        else:
            raise TypeError(
                f"fan-out: component of type "
                f"{type(c).__name__} is neither a "
                f"ContinuousMorphism nor a V-Cat Morphism"
            )
    domain = wrapped_components[0].domain
    for i, c in enumerate(wrapped_components[1:], 1):
        dom_dim = _event_dim(domain)
        c_dim = _event_dim(c.domain)
        if dom_dim != c_dim:
            raise TypeError(
                f"fan-out: component {i} domain dim {c_dim} != component 0 domain dim {dom_dim}"
            )
    codomain = wrapped_components[0].codomain
    for c in wrapped_components[1:]:
        codomain = _combine_spaces(codomain, c.codomain)
    super().__init__(domain, codomain)
    self._components = torch.nn.ModuleList(wrapped_components)
    self._cod_dims = [_event_dim(c.codomain) for c in wrapped_components]

rsample

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

Sample from all components and concatenate outputs.

PARAMETER	DESCRIPTION
x	Input tensor (broadcast to all components). TYPE: Tensor
sample_shape	Additional leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Concatenated outputs from all components.

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample from all components and concatenate outputs.

    Parameters
    ----------
    x : torch.Tensor
        Input tensor (broadcast to all components).
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Concatenated outputs from all components.
    """
    outs = []
    for comp in self._components:
        y = cast(ContinuousMorphism, comp).rsample(x, sample_shape)
        if y.dim() == 1:
            y = y.unsqueeze(-1)
        outs.append(y)
    return torch.cat(outs, dim=-1)

base_dimension

base_dimension(x: Tensor) -> int | None

Sum of the components' coordinate budgets.

Fan-out copies its input to independent components, so their reparameterizations share the input and nothing else.

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

def base_dimension(self, x: torch.Tensor) -> int | None:
    """Sum of the components' coordinate budgets.

    Fan-out copies its input to independent components, so their
    reparameterizations share the input and nothing else.
    """
    dimensions = self._component_dimensions(x)
    if dimensions is None:
        return None
    return sum(dimensions)

push_base

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

Push each component's own coordinate block through it.

The blocks are disjoint and the outputs concatenate along the feature axis, matching the layout rsample and log_prob already use for the fan's codomain.

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

def push_base(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
    """Push each component's own coordinate block through it.

    The blocks are disjoint and the outputs concatenate along the
    feature axis, matching the layout `rsample` and `log_prob`
    already use for the fan's codomain.
    """
    dimensions = self._component_dimensions(x)
    if dimensions is None:
        raise ValueError(
            "FanOutMorphism.push_base: component(s) "
            f"{[type(c).__name__ for c in self._components]!r} "
            "declare no reparameterization, so the fan has none "
            "either."
        )
    outs = []
    offset = 0
    for comp, dimension in zip(self._components, dimensions):
        y = cast(ContinuousMorphism, comp).push_base(
            x, base[:, offset : offset + dimension]
        )
        if y.dim() == 1:
            y = y.unsqueeze(-1)
        outs.append(y)
        offset += dimension
    return torch.cat(outs, dim=-1)

log_prob

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

Log-probability: sum of component log-probs.

PARAMETER	DESCRIPTION
x	Input (same for all components). TYPE: Tensor
y	Concatenated output values. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Sum of log-probabilities. Shape (batch,).

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability: sum of component log-probs.

    Parameters
    ----------
    x : torch.Tensor
        Input (same for all components).
    y : torch.Tensor
        Concatenated output values.

    Returns
    -------
    torch.Tensor
        Sum of log-probabilities. Shape ``(batch,)``.
    """
    lp = torch.zeros(x.shape[0], device=x.device)
    offset = 0
    for comp_mod, d in zip(self._components, self._cod_dims):
        comp = cast(ContinuousMorphism, comp_mod)
        y_slice = y[..., offset : offset + d]
        if _is_discrete(comp.codomain):
            y_slice = y_slice.squeeze(-1).long()
        lp = lp + comp.log_prob(x, y_slice)
        offset += d
    return lp

DiscreteAsContinuous

DiscreteAsContinuous(inner: object)

Bases: ContinuousMorphism

Wrap a discrete Morphism as a ContinuousMorphism.

Enables composition between discrete and continuous morphisms via the >> operator. The wrapped morphism's tensor is used for both log_prob evaluation and sampling.

Note

Sampling from a discrete distribution is NOT reparameterizable. Gradients do not flow through the discrete samples back to the left morphism's parameters. Use score function estimators (REINFORCE) if gradients through discrete choices are needed.

PARAMETER	DESCRIPTION
inner	The discrete morphism to wrap. TYPE: Morphism

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

def __init__(self, inner: object) -> None:
    from quivers.core.morphisms import Morphism

    if not isinstance(inner, Morphism):
        raise TypeError(f"expected a discrete Morphism, got {type(inner).__name__}")
    super().__init__(inner.domain, inner.codomain)
    self._inner = inner
    self._inner_module = inner.module()

log_prob

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

Log-probability from the discrete tensor.

PARAMETER	DESCRIPTION
x	Domain indices. Shape (batch,). TYPE: Tensor
y	Codomain indices. Shape (batch,). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Log-probabilities. Shape (batch,).

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability from the discrete tensor.

    Parameters
    ----------
    x : torch.Tensor
        Domain indices. Shape (batch,).
    y : torch.Tensor
        Codomain indices. Shape (batch,).

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape (batch,).
    """
    t = self._inner.tensor
    probs = t[x.long(), y.long()]
    return torch.log(probs.clamp(min=1e-07))

rsample

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

Sample from the categorical distribution defined by the tensor.

Note: not reparameterizable. Gradients do not flow through the returned samples.

PARAMETER	DESCRIPTION
x	Domain indices. Shape (batch,). TYPE: Tensor
sample_shape	Additional sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Sampled codomain indices. Shape (*sample_shape, batch).

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample from the categorical distribution defined by the tensor.

    Note: not reparameterizable. Gradients do not flow through
    the returned samples.

    Parameters
    ----------
    x : torch.Tensor
        Domain indices. Shape (batch,).
    sample_shape : torch.Size
        Additional sample dimensions.

    Returns
    -------
    torch.Tensor
        Sampled codomain indices. Shape (*sample_shape, batch).
    """
    t = self._inner.tensor
    probs = t[x.long()]
    n_samples = (
        int(torch.Size(sample_shape).numel()) if len(sample_shape) > 0 else 1
    )
    samples = torch.multinomial(probs, n_samples, replacement=True)
    if len(sample_shape) == 0:
        return samples.squeeze(-1)
    else:
        return samples.T.reshape(*sample_shape, -1)

dimension_probe

dimension_probe(x: Tensor) -> Tensor

A one-row slice of x, enough to settle coordinate counts.

A morphism's base_dimension depends on trailing event extents, not the number of rows.

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

def dimension_probe(x: torch.Tensor) -> torch.Tensor:
    """A one-row slice of ``x``, enough to settle coordinate counts.

    A morphism's
    [`base_dimension`][quivers.continuous.morphisms.ContinuousMorphism.base_dimension]
    depends on trailing event extents, not the number of rows.
    """
    return x[:1]

sobol_normal_points

sobol_normal_points(dimension: int, count: int, device: device, dtype: dtype) -> Tensor

A deterministic standard-normal point set of shape (n, dimension).

Push an unscrambled Sobol point set through the standard-normal quantile function. The result is deterministic. The implementation skips the Sobol origin, clamps quantile inputs away from 0 and 1, and rounds count up to a power of two.

PARAMETER	DESCRIPTION
dimension	Number of coordinates per point. Zero yields an empty (n, 0) tensor, which is what a deterministic map consumes. TYPE: int
count	Requested point count; rounded up to a power of two. TYPE: int
device	Device to place the result on. TYPE: device
dtype	Floating dtype of the result. TYPE: dtype

RETURNS	DESCRIPTION
Tensor	Shape (n, dimension).

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

def sobol_normal_points(
    dimension: int,
    count: int,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    """A deterministic standard-normal point set of shape ``(n, dimension)``.

    Push an unscrambled Sobol point set through the standard-normal
    quantile function. The result is deterministic. The implementation
    skips the Sobol origin, clamps quantile inputs away from 0 and 1,
    and rounds ``count`` up to a power of two.

    Parameters
    ----------
    dimension : int
        Number of coordinates per point. Zero yields an empty
        ``(n, 0)`` tensor, which is what a deterministic map consumes.
    count : int
        Requested point count; rounded up to a power of two.
    device : torch.device
        Device to place the result on.
    dtype : torch.dtype
        Floating dtype of the result.

    Returns
    -------
    torch.Tensor
        Shape ``(n, dimension)``.
    """
    n = _next_power_of_two(count)
    if dimension == 0:
        return torch.zeros(n, 0, device=device, dtype=dtype)
    engine = torch.quasirandom.SobolEngine(dimension=dimension, scramble=False)
    unit = engine.draw(n + 1, dtype=torch.float64)[1:]
    unit = unit.clamp(min=_QUANTILE_EPS, max=1.0 - _QUANTILE_EPS)
    return torch.special.ndtri(unit).to(device=device, dtype=dtype)

chain_dimensions

chain_dimensions(factors: 'collections.abc.Sequence[ContinuousMorphism]', x: Tensor) -> list[int] | None

Per-factor base-coordinate counts along a chain, or None.

A factor's count can depend on the shape of what reaches it, so the chain is walked once with the coordinates held at zero. That pushes each kernel's median forward, which costs a forward pass and settles the shapes without consuming a point set the caller has not built yet.

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

def chain_dimensions(
    factors: "collections.abc.Sequence[ContinuousMorphism]", x: torch.Tensor
) -> list[int] | None:
    """Per-factor base-coordinate counts along a chain, or None.

    A factor's count can depend on the shape of what reaches it, so
    the chain is walked once with the coordinates held at zero. That
    pushes each kernel's median forward, which costs a forward pass
    and settles the shapes without consuming a point set the caller
    has not built yet.
    """
    dimensions: list[int] = []
    probe = dimension_probe(x)
    dtype = x.dtype if x.is_floating_point() else torch.get_default_dtype()
    for factor in factors:
        dimension = factor.base_dimension(probe)
        if dimension is None:
            return None
        dimensions.append(dimension)
        zeros = torch.zeros(probe.shape[0], dimension, device=probe.device, dtype=dtype)
        probe = factor.push_base(probe, zeros)
    return dimensions

chain_push_base

chain_push_base(factors: 'collections.abc.Sequence[ContinuousMorphism]', x: Tensor, base: Tensor, dimensions: list[int]) -> Tensor

Thread base coordinates through a chain, one block per factor.

Each factor consumes its own contiguous block, so no two factors read the same coordinate and the composite map is the pushforward of a single point set through the whole chain rather than per-factor rules glued together by index. Sharing coordinates across factors would resolve some directions of the joint twice and leave others unexplored.

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

def chain_push_base(
    factors: "collections.abc.Sequence[ContinuousMorphism]",
    x: torch.Tensor,
    base: torch.Tensor,
    dimensions: list[int],
) -> torch.Tensor:
    """Thread base coordinates through a chain, one block per factor.

    Each factor consumes its own contiguous block, so no two factors
    read the same coordinate and the composite map is the pushforward
    of a single point set through the whole chain rather than
    per-factor rules glued together by index. Sharing coordinates
    across factors would resolve some directions of the joint twice
    and leave others unexplored.
    """
    offset = 0
    current = x
    for factor, dimension in zip(factors, dimensions):
        current = factor.push_base(current, base[:, offset : offset + dimension])
        offset += dimension
    return current

chain_marginal_quadrature

chain_marginal_quadrature(factors: 'collections.abc.Sequence[ContinuousMorphism]', x: Tensor, count: int) -> tuple[Tensor, Tensor] | None

A deterministic rule for the law a whole chain induces on its end.

Push one point set of dimension sum(chain_dimensions(...)) through every factor. The rule returns n terminal nodes without multiplying the node count at each link.

PARAMETER	DESCRIPTION
factors	The chain, in application order. A single-element sequence defers to that morphism's own rule. TYPE: Sequence[ContinuousMorphism]
x	Conditioning inputs. Shape (batch, *domain). TYPE: Tensor
count	Requested node count; rounded up to a power of two. TYPE: int

RETURNS	DESCRIPTION
tuple[Tensor, Tensor] or None	Nodes of shape (n, batch, *event) and log-weights of shape (n,), or None when any factor has no reparameterization.

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

def chain_marginal_quadrature(
    factors: "collections.abc.Sequence[ContinuousMorphism]",
    x: torch.Tensor,
    count: int,
) -> tuple[torch.Tensor, torch.Tensor] | None:
    """A deterministic rule for the law a whole chain induces on its end.

    Push one point set of dimension ``sum(chain_dimensions(...))``
    through every factor. The rule returns ``n`` terminal nodes without
    multiplying the node count at each link.

    Parameters
    ----------
    factors : Sequence[ContinuousMorphism]
        The chain, in application order. A single-element sequence
        defers to that morphism's own rule.
    x : torch.Tensor
        Conditioning inputs. Shape ``(batch, *domain)``.
    count : int
        Requested node count; rounded up to a power of two.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor] or None
        Nodes of shape ``(n, batch, *event)`` and log-weights of
        shape ``(n,)``, or ``None`` when any factor has no
        reparameterization.
    """
    if not factors:
        raise ValueError(
            "chain_marginal_quadrature: an empty chain induces no law; "
            "pass at least one factor."
        )
    if len(factors) == 1:
        return factors[0].marginal_quadrature(x, count)
    dimensions = chain_dimensions(factors, x)
    if dimensions is None:
        return None
    total = sum(dimensions)
    batch = x.shape[0]
    dtype = x.dtype if x.is_floating_point() else torch.get_default_dtype()
    base = sobol_normal_points(total, count, x.device, dtype)
    n = base.shape[0]
    x_rows = x.unsqueeze(0).expand(n, *x.shape).reshape(n * batch, *x.shape[1:])
    base_rows = base.unsqueeze(1).expand(n, batch, total).reshape(n * batch, total)
    pushed = chain_push_base(factors, x_rows, base_rows, dimensions)
    nodes = pushed.reshape(n, batch, *pushed.shape[1:])
    log_weights = torch.full(
        (n,), -math.log(float(n)), device=nodes.device, dtype=nodes.dtype
    )
    return nodes, log_weights