Plates

Indexed draws and observations. A PlateDraw draws one value per element of a finite set, as sample v : A <- F(...) does in source; a VectorisedObserve scores a column of responses at once, as an indexed observe does. The marginalization helpers integrate a finite latent out of a plate of scores.

plate

Hierarchical Bayesian modelling primitives.

Adds the operations needed to express models of the kind found in hierarchical Bayesian regression / latent-class analysis (the canonical Stan workhorse): finite-domain-indexed draws (plates), vectorized observations with gather indexing, LKJ priors on correlation matrices, ordinal monotone splines via cumulative sum, and generic distribution truncation. Each primitive is declared with its categorical denotation in :math:\mathbf{Kern}; the runtime is a straight realisation of those denotations.

Categorical foundations

• Plate (indexed draw). Given a finite index set :math:A and a parameterized family :math:F : \Theta \to \mathcal{G}(B), a plate draw declares the kernel

.. math::

  v : A \to B,\quad v(a) \sim F(\theta(a))

equivalently a single morphism into the function space :math:\mathbf{1} \to \mathcal{G}(B^A) factoring as the independent product :math:\prod_{a \in A} F(\theta(a)). Under the natural isomorphism :math:\mathbf{Kern}(\mathbf{1}, B^A) \cong \mathbf{Kern}(A, B) the plate variable IS a :math:\mathbf{Kern}-morphism :math:A \to B; in the variational backend it is held as a tensor of shape (|A|, *B.shape) whose prior contribution to the ELBO is the per-row log-density.

• Gather. Given a finite fibration :math:\iota : N \to A (a per-observation grouping assignment) and a plate variable :math:v : A \to B, the gathered morphism is the pullback :math:\iota^* v = v \circ \iota : N \to B. Categorically a textbook reindexing morphism in :math:\mathbf{Kern}; runtime is v[indices] along the leading axis.

• Vectorised observe. A batched observation step

.. math::

  \mathcal{S}\llbracket\,\mathsf{observe}\ r[n] \sim
  F(\theta[n])\ \mathsf{for}\ n\in N\,\rrbracket :
  \Phi \to \mathcal{G}_{\le 1}(\Phi)

has score :math:\prod_{n \in N} p_F(r_{\text{obs}}(n);\, \theta(n,\phi)). The single-observation form is the |N| = 1 special case; mass-correctness follows from the Cho-Jacobs Markov-with-conditioning calculus.

• Marginalise (program-level). Given a previously-drawn discrete latent :math:c : \Phi \to \mathcal{G}(C), the marginalisation step is the pushforward through the projection :math:\pi_{\Phi \setminus C} : \Phi \times C \to \Phi:

.. math::

  \mathsf{marg}(c) = \mathcal{G}(\pi_{\Phi \setminus C})
                     \circ \mathcal{S}\llbracket\mathsf{draw}\ c\rrbracket

numerically realised as :math:\log \sum_c \exp(\cdot) over the :math:C axis of the per-class log-likelihoods.

• cumsum. The deterministic morphism :math:\mathrm{cumsum} : \mathrm{Euclidean}(K) \to \mathrm{Euclidean}(K), :math:(x_1, \dots, x_K) \mapsto (x_1, x_1 + x_2, \dots, \sum_i x_i). Used for the standard monotonic-spline parameterization of ordinal coefficients.

• LKJ correlation prior. The Lewandowski-Kurowicka-Joe distribution :math:\mathrm{LKJ}(K, \eta) on the manifold :math:\mathrm{Corr}_K of :math:K \times K correlation matrices, parameterized via the Cholesky factor for numerical stability. The accompanying cholesky_quad_form deterministic morphism reconstructs a full covariance :math:\Sigma = \mathrm{diag}(s)\, R\, \mathrm{diag}(s) from a correlation :math:R and a per-component scale :math:s \in \mathrm{PositiveReals}(K).

• Truncated. Given a base family :math:F and an interval :math:[a, b], the truncated family :math:F_{|[a, b]} has density :math:p_F(x) / (F_{\text{cdf}}(b) - F_{\text{cdf}}(a)) on :math:[a, b] and zero elsewhere. Realised via inverse-CDF sampling when available; rejection otherwise.

References

• Cho, K. and Jacobs, B. (2019). Disintegration and Bayesian inversion via string diagrams. Mathematical Structures in Computer Science 29(7), 938–971. doi:10.1017/S0960129518000488.
• Lewandowski, D., Kurowicka, D. and Joe, H. (2009). Generating random correlation matrices based on vines and extended onion method. Journal of Multivariate Analysis 100(9), 1989–2001. doi:10.1016/j.jmva.2009.04.008.

PlateDraw

PlateDraw(index_size: int, family: ContinuousMorphism, domain: AnySpace | None = None)

Bases: ContinuousMorphism

A finite-domain-indexed draw, as a Kern-morphism A → B.

Concretely: v : A → B ~ F(theta) becomes a tensor of shape (|A|, *B.shape) whose a-th row is an independent :math:F(\theta_a)-distributed random variable. The variational posterior factorizes across rows by default; the prior's ELBO contribution is :math:\sum_a \log p_F(v_a; \theta_a).

Categorically: by the natural isomorphism :math:\mathbf{Kern}(\mathbf{1}, B^A) \cong \mathbf{Kern}(A, B), the plate variable IS a Kern-morphism :math:A \to B. The PlateDraw is realised as a ContinuousMorphism whose codomain is the flat product-space of index_size copies of the per-row family's codomain.

PARAMETER	DESCRIPTION
index_size	Cardinality :math:|A|. TYPE: int
family	Per-row distribution family. TYPE: ContinuousMorphism
domain	The program's input space (broadcast conditioning). TYPE: AnySpace DEFAULT: None

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

def __init__(
    self,
    index_size: int,
    family: ContinuousMorphism,
    domain: AnySpace | None = None,
) -> None:
    # A matrix-valued space (a Cholesky factor, a covariance) keeps
    # `dim` as its matrix size and states its flat element count as
    # `shape`; other continuous spaces state only `dim`.
    per_row_shape: tuple[int, ...]
    if isinstance(family.codomain, _MATRIX_SPACES):
        per_row_shape = family.codomain.shape
        per_row_dim = math.prod(per_row_shape)
    elif hasattr(family.codomain, "dim"):
        per_row_dim = int(family.codomain.dim)
        per_row_shape = (per_row_dim,)
    else:
        per_row_shape = tuple(family.codomain.shape)
        per_row_dim = (
            int(torch.tensor(per_row_shape).prod().item()) if per_row_shape else 1
        )
    # `plate_rows` keeps the row structure of the flattened
    # codomain recoverable: a consumer that reads the variable as
    # a *parameter* needs the per-row width, since the plate axis
    # is a batch axis in that position rather than an event axis.
    flat_codomain = Euclidean(
        name=f"plate({index_size}x{family.codomain!s})",
        dim=index_size * per_row_dim,
        plate_rows=index_size,
    )
    actual_domain = domain if domain is not None else family.domain
    super().__init__(actual_domain, flat_codomain)
    self._index_size = index_size
    self._family = family
    self._per_row_shape = per_row_shape
    self._per_row_dim = per_row_dim
    # Width of the conditioning input a single row of the family
    # consumes. Fixed families ignore it (they read only the row
    # count); conditional families feed it through their parameter
    # network, so it must equal the domain's feature width.
    dom = actual_domain
    if hasattr(dom, "dim"):
        self._domain_width = max(1, int(dom.dim))
    elif hasattr(dom, "shape"):
        shape = tuple(dom.shape)
        self._domain_width = int(torch.tensor(shape).prod().item()) if shape else 1
    else:
        self._domain_width = 1

rsample

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

Reparameterized sample from the plate's prior.

A plate draw is batch-invariant: the latent vector is a global model quantity shared across every row of an observed plate. The returned tensor has shape (|A|, *B.shape) regardless of the program input's leading batch dimension, with each row drawn independently from the per-row family F at its declared parameters. x supplies only the device; it is optional so batch-invariant callers such as gather need not fabricate one.

This is the standard Pyro / NumPyro semantic: a sample inside a plate("subj", n_subj) context is one (n_subj,) vector, not a per-particle (batch, n_subj) replication. The gather arr[idx] along the plate axis then composes cleanly with per-row observed-plate axes downstream.

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

def rsample(
    self,
    x: torch.Tensor | None = None,
    sample_shape: torch.Size = torch.Size(),
) -> torch.Tensor:
    """Reparameterized sample from the plate's prior.

    A plate draw is *batch-invariant*: the latent vector is a
    global model quantity shared across every row of an observed
    plate. The returned tensor has shape ``(|A|, *B.shape)``
    regardless of the program input's leading batch dimension,
    with each row drawn independently from the per-row family
    ``F`` at its declared parameters. ``x`` supplies only the
    device; it is optional so batch-invariant callers such as
    `gather` need not fabricate one.

    This is the standard Pyro / NumPyro semantic: a sample inside
    a ``plate("subj", n_subj)`` context is one ``(n_subj,)``
    vector, *not* a per-particle (batch, n_subj) replication. The
    gather ``arr[idx]`` along the plate axis then composes
    cleanly with per-row observed-plate axes downstream.
    """
    del sample_shape  # plate latents are batch-invariant
    device = x.device if x is not None else None
    # Draw ``|A|`` independent rows by sampling the per-row family
    # with a leading sample axis of length ``|A|`` from a single
    # conditioning row. Threading the row count through
    # ``sample_shape`` rather than the input's batch axis keeps the
    # draw genuinely independent even for families whose fixed
    # parameterization advertises a scalar batch shape (a
    # positive-support family realised as a restricted base
    # measure, say); those ignore the input's leading axis but
    # still honour the explicit sample shape.
    row_input = self._row_conditioning(x, 1, device)
    base = self._family.rsample(
        row_input, sample_shape=torch.Size((self._index_size,))
    )
    expected = self._index_size * self._per_row_dim
    if base.numel() != expected:
        raise ValueError(
            f"plate rsample: per-row family produced {base.numel()} "
            f"elements for {self._index_size} rows of per-row width "
            f"{self._per_row_dim}; expected {expected}"
        )
    return self._to_structured(base)

canonical_latent

canonical_latent(value: Tensor) -> Tensor

Give a clamped plate value its structured (|A|, *B.shape).

A plate latent supplied as a conditioning / observed value may arrive flattened to (|A| * prod(B.shape),) (the shape a test point or host-data channel produces when it flattens a matrix-valued latent). Restoring the structured shape here lets a gather v[idx] over the plate axis yield (N, *B.shape) with the per-row event axis intact.

A value whose element count is not exactly |A| * prod(B) is a genuinely batched latent (a leading parameter-sample axis); it is returned unchanged so the batched scoring path handles it.

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

def canonical_latent(self, value: torch.Tensor) -> torch.Tensor:
    """Give a clamped plate value its structured ``(|A|, *B.shape)``.

    A plate latent supplied as a conditioning / observed value may
    arrive flattened to ``(|A| * prod(B.shape),)`` (the shape a
    test point or host-data channel produces when it flattens a
    matrix-valued latent). Restoring the structured shape here lets
    a gather ``v[idx]`` over the plate axis yield ``(N, *B.shape)``
    with the per-row event axis intact.

    A value whose element count is not exactly ``|A| * prod(B)`` is
    a genuinely batched latent (a leading parameter-sample axis);
    it is returned unchanged so the batched scoring path handles
    it.
    """
    if value.numel() != self._index_size * self._per_row_dim:
        return value
    return self._to_structured(value)

log_prob

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

Prior log-density of the plate sample.

The plate variable v : A -> B scores as the independent product prod_a p_F(v_a) of the per-row family F; the contribution is sum_a log p_F(v_a).

Two layouts are accepted. The natural plate latent carries a leading axis of length |A| (or is a flat vector whose element count is |A| * prod(event)): its rows are scored under F and reduced to a scalar, wrapped in a length-1 tensor so the downstream log-joint accumulator sees a uniform 1-d shape it can broadcast against the response plate. A batched value carries a genuine leading batch axis distinct from |A|; each batch element is scored independently and a (batch,) vector of per-element densities is returned.

The per-row event size is read from the value's element count rather than solely from the family's advertised codomain, so a scalar family carrying an explicit event axis (over=...) is scored over the full event without loss.

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Prior log-density of the plate sample.

    The plate variable ``v : A -> B`` scores as the independent
    product ``prod_a p_F(v_a)`` of the per-row family ``F``; the
    contribution is ``sum_a log p_F(v_a)``.

    Two layouts are accepted. The natural *plate latent* carries a
    leading axis of length ``|A|`` (or is a flat vector whose
    element count is ``|A| * prod(event)``): its rows are scored
    under ``F`` and reduced to a scalar, wrapped in a length-1
    tensor so the downstream log-joint accumulator sees a uniform
    1-d shape it can broadcast against the response plate. A
    *batched* value carries a genuine leading batch axis distinct
    from ``|A|``; each batch element is scored independently and a
    ``(batch,)`` vector of per-element densities is returned.

    The per-row event size is read from the value's element count
    rather than solely from the family's advertised codomain, so a
    scalar family carrying an explicit event axis (``over=...``)
    is scored over the full event without loss.
    """
    index_size = self._index_size
    per_row_family = self._per_row_dim
    batched = y.dim() >= 2 and y.shape[0] != index_size
    if not batched:
        total = y.numel()
        if total % index_size != 0:
            raise ValueError(
                f"plate latent has {total} elements, not divisible "
                f"by the plate index size {index_size}"
            )
        row_numel = total // index_size
        event_shape = (
            self._per_row_shape if row_numel == per_row_family else (row_numel,)
        )
        sample = y.reshape(index_size, *event_shape)
        row_input = self._row_conditioning(x, index_size, y.device)
        per_row_lp = self._family.log_prob(row_input, sample)
        return per_row_lp.reshape(-1).sum().unsqueeze(0)
    batch = y.shape[0]
    per_batch = y.numel() // batch
    if per_batch % index_size != 0:
        raise ValueError(
            f"batched plate latent has {per_batch} elements per "
            f"batch row, not divisible by the plate index size "
            f"{index_size}"
        )
    row_numel = per_batch // index_size
    event_shape = (
        self._per_row_shape if row_numel == per_row_family else (row_numel,)
    )
    flat = y.reshape(batch * index_size, *event_shape)
    row_input = self._row_conditioning(x, batch * index_size, y.device)
    per_row_lp = self._family.log_prob(row_input, flat)
    return per_row_lp.reshape(batch, index_size).sum(dim=-1)

gather

gather(indices: Tensor) -> Tensor

Pullback v[indices] along a finite fibration.

indices is a long-tensor of shape (N,) with entries in [0, |A|); returns a tensor of shape (N, *B.shape).

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

def gather(self, indices: torch.Tensor) -> torch.Tensor:
    """Pullback ``v[indices]`` along a finite fibration.

    ``indices`` is a long-tensor of shape ``(N,)`` with entries
    in ``[0, |A|)``; returns a tensor of shape
    ``(N, *B.shape)``.
    """
    return self.rsample()[indices]

VectorisedObserve

VectorisedObserve(family: ContinuousMorphism, response: Tensor)

Bases: ContinuousMorphism

A batched observation step accumulating per-row log-likelihoods.

Categorically, the batched-likelihood kernel :math:\Phi \to \mathcal{G}_{\le 1}(\Phi) whose score is :math:\prod_{n \in N} p_F(r_{\text{obs}}(n);\, \theta(n,\phi)). Realised as a ContinuousMorphism whose domain is the parameter-input space (the morphism conditions on θ) and whose codomain is the per-observation response space — so the existing MonadicProgram observe-step machinery treats it as an observed site and threads the score through log_joint via the usual morph.log_prob(theta, response) call, with log_prob here summing over the leading index axis.

The observed response tensor is registered as a buffer so the parent program's optimizer tracks it and the runtime never has to thread it through observations=....

PARAMETER	DESCRIPTION
family	The per-observation distribution family. TYPE: ContinuousMorphism
response	Observed values r_obs of shape (N, *codom.shape). A one-dimensional continuous codomain also accepts (N,), which gains its event axis in _align_response; a discrete codomain takes (N,) indices and has no event axis. TYPE: Tensor

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

def __init__(self, family: ContinuousMorphism, response: torch.Tensor) -> None:
    super().__init__(family.domain, family.codomain)
    self._family = family
    self.register_buffer("_response", response.detach())

support property

support

Inherit the underlying per-row family's support so guides and validators see the correct constraint for the observation.

rsample

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

Sample the per-observation family at the supplied θ.

x is the θ-tensor (one row per observation index); the result is the per-observation response sample. Used in prior-predictive simulation; never called during inference when the response is observed.

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample the per-observation family at the supplied θ.

    ``x`` is the θ-tensor (one row per observation index); the
    result is the per-observation response sample. Used in
    prior-predictive simulation; never called during inference
    when the response is observed.
    """
    return self._family.rsample(x, sample_shape)

log_prob

log_prob(x: Tensor, y: Tensor | None = None) -> Tensor

Sum of per-observation log-densities.

y defaults to the registered response buffer; passing a different value (e.g. a clamped observation) is supported for fast prior-predictive checks.

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor | None = None) -> torch.Tensor:
    """Sum of per-observation log-densities.

    ``y`` defaults to the registered response buffer; passing a
    different value (e.g. a clamped observation) is supported
    for fast prior-predictive checks.
    """
    target = y if y is not None else cast("torch.Tensor", self._response)
    return self._family.log_prob(x, self._align_response(target)).sum()

log_likelihood

log_likelihood(theta: Tensor) -> Tensor

Alias for log_prob(theta); preserved for the Python builder API.

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

def log_likelihood(self, theta: torch.Tensor) -> torch.Tensor:
    """Alias for ``log_prob(theta)``; preserved for the Python
    builder API."""
    return self.log_prob(theta)

marginalize_categorical

marginalize_categorical(log_probs_per_class: Tensor) -> Tensor

Program-level marginalisation over a discrete latent class.

Given per-class log-likelihoods of shape (N, K) (one row per observation, one column per latent-class assignment), returns the marginalised log-likelihood of shape (N,):

.. math::

\log p(r_n) = \log \sum_{k=1}^{K} p(c_n = k)\, p(r_n | c_n = k).

Realises the program-level marginalize c step as the pushforward through :math:\pi_{\Phi \setminus C}.

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

def marginalize_categorical(log_probs_per_class: torch.Tensor) -> torch.Tensor:
    """Program-level marginalisation over a discrete latent class.

    Given per-class log-likelihoods of shape ``(N, K)`` (one row per
    observation, one column per latent-class assignment), returns
    the marginalised log-likelihood of shape ``(N,)``:

    .. math::

        \\log p(r_n) = \\log \\sum_{k=1}^{K} p(c_n = k)\\, p(r_n | c_n = k).

    Realises the program-level ``marginalize c`` step as the
    pushforward through :math:`\\pi_{\\Phi \\setminus C}`.
    """
    return torch.logsumexp(log_probs_per_class, dim=-1)

marginalize_grouped

marginalize_grouped(log_likelihood_per_row_per_class: Tensor | list[Tensor] | tuple[Tensor, ...], group_index: Tensor | tuple[Tensor, ...] | list[Tensor | tuple[Tensor, ...]], log_prior_per_group_per_class: Tensor, num_groups: int | tuple[int, ...], *, reduction: str = 'logsumexp', per_group: bool = False) -> Tensor

Per-group marginalisation over a discrete latent class.

Two surface shapes are accepted, the second a strict generalisation of the first.

Single response axis. Given a per-(response, class) log-likelihood tensor of shape (N, K) and a fibration :math:r : \text{Resp} \to G (or a tuple of co-indexed fibrations into a product grouping plate :math:G_1 \times \dots \times G_r), return the scalar

.. math::

\sum_{g \in G}\, \rho\!\left[
    \log \pi(g, \cdot) + \sum_{n:\, r(n) = g} \ell(n, \cdot)
\right].

Multiple heterogeneous response axes sharing a per-group class indicator. Given M per-axis log-likelihood tensors :math:\ell_m \in \mathbb{R}^{N_m \times K} and a parallel list of fibrations :math:r_m : \text{Resp}_m \to G, return

.. math::

\sum_{g \in G}\, \rho\!\left[
    \log \pi(g, \cdot)
    + \sum_{m=1}^{M} \sum_{n:\, r_m(n) = g} \ell_m(n, \cdot)
\right].

Each per-axis log-likelihood is scatter-added into the same (|G|, K) accumulator before the reduction. Categorically this is the right Kan extension along the coproduct fibration :math:\coprod_m r_m : \coprod_m \text{Resp}_m \to G. The single-axis form is the unary case (M = 1).

The reduction :math:\rho over the class axis is selected by reduction: logsumexp for the canonical mixture marginalisation, sum for joint scoring (no marginalisation), mean for the symmetric average.

PARAMETER	DESCRIPTION
log_likelihood_per_row_per_class	Single tensor: a per-(response, class) log-likelihood of shape (N, K), used in the single-axis form. List or tuple of tensors: per-axis log-likelihoods of shapes (N_m, K), used in the multi-axis form. The class axis K must agree across all entries. TYPE: torch.Tensor or list/tuple of torch.Tensor
group_index	Single-axis form: a long tensor of shape (N,) (single fibration) or a tuple of co-indexed long tensors (product fibration). Multi-axis form: a list of length M, parallel to the log-likelihood list; each entry is itself either a single tensor or a product-fibration tuple, exactly as in the single-axis form. TYPE: torch.Tensor, tuple of torch.Tensor, or list of either
log_prior_per_group_per_class	Per-(group, class) log-prior of shape (num_groups, K), (K,) (broadcast over the group axis), or, for product fibrations, (num_groups_1, ..., num_groups_r, K). TYPE: Tensor
num_groups	Cardinality of the (product) group plate. TYPE: int or tuple[int, ...]
reduction	Per-group reduction over the class axis. logsumexp is the canonical mixture marginalisation; sum joint-scores without marginalising the class; mean averages symmetrically. Default logsumexp. TYPE: ('logsumexp', 'sum', 'mean') DEFAULT: "logsumexp"
per_group	Whether to return the reduced per-group aggregates (|G|, *extra) (flat over a product plate) instead of their sum over the group axis. A block nested inside another grouped block contributes one aggregate per position of its group to the enclosing accumulator, so the enclosing block consumes this form. TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
Tensor	Scalar program-level log-density contribution, or the per-group aggregates when per_group is set.

RAISES	DESCRIPTION
ValueError	If reduction is not one of the three reductions, if the fibration and log-likelihood shapes disagree, if a prior has a shape that fits neither the group plate nor the class axis, or if per_group is requested for a log-likelihood with no row axis (there is no group axis to keep).

Notes

Edge cases:

• K == 1: the log-sum-exp collapses to the body's log-likelihood plus the constant log-prior.
• Identity fibration (group_index = arange(N) and num_groups == N): each row is its own group, recovering the per-row mixture.
• Empty fibre: contributes a multiplicative identity to the product (log 1 = 0), so an unused group adds zero per class before the reduction.
• Multi-axis with a single entry (M = 1): exactly equivalent to the single-axis form, modulo unboxing the singleton list.

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

def marginalize_grouped(
    log_likelihood_per_row_per_class: (
        torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...]
    ),
    group_index: (
        torch.Tensor
        | tuple[torch.Tensor, ...]
        | list[torch.Tensor | tuple[torch.Tensor, ...]]
    ),
    log_prior_per_group_per_class: torch.Tensor,
    num_groups: int | tuple[int, ...],
    *,
    reduction: str = "logsumexp",
    per_group: bool = False,
) -> torch.Tensor:
    """Per-group marginalisation over a discrete latent class.

    Two surface shapes are accepted, the second a strict
    generalisation of the first.

    **Single response axis.** Given a per-(response, class)
    log-likelihood tensor of shape ``(N, K)`` and a fibration
    :math:`r : \\text{Resp} \\to G` (or a tuple of co-indexed
    fibrations into a product grouping plate
    :math:`G_1 \\times \\dots \\times G_r`), return the scalar

    .. math::

        \\sum_{g \\in G}\\, \\rho\\!\\left[
            \\log \\pi(g, \\cdot) + \\sum_{n:\\, r(n) = g} \\ell(n, \\cdot)
        \\right].

    **Multiple heterogeneous response axes sharing a per-group
    class indicator.** Given ``M`` per-axis log-likelihood tensors
    :math:`\\ell_m \\in \\mathbb{R}^{N_m \\times K}` and a parallel
    list of fibrations :math:`r_m : \\text{Resp}_m \\to G`, return

    .. math::

        \\sum_{g \\in G}\\, \\rho\\!\\left[
            \\log \\pi(g, \\cdot)
            + \\sum_{m=1}^{M} \\sum_{n:\\, r_m(n) = g} \\ell_m(n, \\cdot)
        \\right].

    Each per-axis log-likelihood is scatter-added into the same
    ``(|G|, K)`` accumulator before the reduction. Categorically
    this is the right Kan extension along the coproduct fibration
    :math:`\\coprod_m r_m : \\coprod_m \\text{Resp}_m \\to G`. The
    single-axis form is the unary case (``M = 1``).

    The reduction :math:`\\rho` over the class axis is selected by
    ``reduction``: ``logsumexp`` for the canonical mixture
    marginalisation, ``sum`` for joint scoring (no
    marginalisation), ``mean`` for the symmetric average.

    Parameters
    ----------
    log_likelihood_per_row_per_class : torch.Tensor or list/tuple of torch.Tensor
        Single tensor: a per-(response, class) log-likelihood of
        shape ``(N, K)``, used in the single-axis form. List or
        tuple of tensors: per-axis log-likelihoods of shapes
        ``(N_m, K)``, used in the multi-axis form. The class axis
        ``K`` must agree across all entries.
    group_index : torch.Tensor, tuple of torch.Tensor, or list of either
        Single-axis form: a long tensor of shape ``(N,)`` (single
        fibration) or a tuple of co-indexed long tensors (product
        fibration). Multi-axis form: a list of length ``M``,
        parallel to the log-likelihood list; each entry is itself
        either a single tensor or a product-fibration tuple, exactly
        as in the single-axis form.
    log_prior_per_group_per_class : torch.Tensor
        Per-(group, class) log-prior of shape ``(num_groups, K)``,
        ``(K,)`` (broadcast over the group axis), or, for product
        fibrations, ``(num_groups_1, ..., num_groups_r, K)``.
    num_groups : int or tuple[int, ...]
        Cardinality of the (product) group plate.
    reduction : {"logsumexp", "sum", "mean"}
        Per-group reduction over the class axis. ``logsumexp`` is
        the canonical mixture marginalisation; ``sum`` joint-scores
        without marginalising the class; ``mean`` averages
        symmetrically. Default ``logsumexp``.
    per_group : bool
        Whether to return the reduced per-group aggregates
        ``(|G|, *extra)`` (flat over a product plate) instead of
        their sum over the group axis. A block nested inside
        another grouped block contributes one aggregate per
        position of its group to the enclosing accumulator, so
        the enclosing block consumes this form.

    Returns
    -------
    torch.Tensor
        Scalar program-level log-density contribution, or the
        per-group aggregates when ``per_group`` is set.

    Raises
    ------
    ValueError
        If ``reduction`` is not one of the three reductions, if the
        fibration and log-likelihood shapes disagree, if a prior has
        a shape that fits neither the group plate nor the class
        axis, or if ``per_group`` is requested for a log-likelihood
        with no row axis (there is no group axis to keep).

    Notes
    -----
    Edge cases:

    * ``K == 1``: the log-sum-exp collapses to the body's
      log-likelihood plus the constant log-prior.
    * Identity fibration (``group_index = arange(N)`` and
      ``num_groups == N``): each row is its own group, recovering
      the per-row mixture.
    * Empty fibre: contributes a multiplicative identity to the
      product (``log 1 = 0``), so an unused group adds zero per
      class before the reduction.
    * Multi-axis with a single entry (``M = 1``): exactly
      equivalent to the single-axis form, modulo unboxing the
      singleton list.
    """
    if isinstance(log_likelihood_per_row_per_class, (list, tuple)):
        if not isinstance(group_index, list):
            raise ValueError(
                "marginalize_grouped: multi-axis log-likelihood "
                "requires a parallel list of group_index entries; "
                f"got {type(group_index).__name__}"
            )
        if len(log_likelihood_per_row_per_class) != len(group_index):
            raise ValueError(
                "marginalize_grouped: log-likelihood list has "
                f"{len(log_likelihood_per_row_per_class)} entries "
                f"but group_index list has {len(group_index)}"
            )
        return _marginalize_grouped_multi(
            tuple(log_likelihood_per_row_per_class),
            tuple(group_index),
            log_prior_per_group_per_class,
            num_groups,
            reduction=reduction,
            per_group=per_group,
        )
    if reduction not in ("logsumexp", "sum", "mean"):
        raise ValueError(
            f"reduction must be one of 'logsumexp', 'sum', 'mean'; got {reduction!r}"
        )
    if log_likelihood_per_row_per_class.dim() < 1:
        raise ValueError(
            "log_likelihood_per_row_per_class must have at least one "
            "axis (the class axis); got "
            f"{log_likelihood_per_row_per_class.shape}"
        )
    # Two operating modes:
    #
    # 1. ``ll`` has a leading row axis (shape ``(N, *extra, K)``):
    #    scatter-add along the fibration to obtain
    #    ``(G, *extra, K)``, apply the per-cell prior, reduce over
    #    the class axis, and sum over groups. The shape of the
    #    return value is ``(*extra,)`` — a scalar if there are no
    #    extra (outer-class) axes, otherwise a tensor whose axes
    #    are the outer-block class axes still in scope.
    #
    # 2. ``ll`` has no row axis (shape ``(*extra, K)``): this
    #    occurs for *intermediate* levels of a nested marginalize
    #    stack, where the innermost level has already integrated
    #    the per-row contributions. There is no scatter step; the
    #    prior is added across the extra axes and the class axis is
    #    reduced. The return shape is ``(*extra_except_class_axis,)``.
    # The mode is detected from the shape of ``group_index``.
    flat_group_index_shape: tuple[int, ...]
    if isinstance(group_index, tuple):
        flat_group_index_shape = group_index[0].shape
    else:
        flat_group_index_shape = group_index.shape
    if flat_group_index_shape:
        n_idx = flat_group_index_shape[0]
        # If the user supplied a fibration of length N > 0, the ll
        # MUST carry a leading axis of length N. A length-N fibration
        # paired with an ll whose leading axis differs is a shape
        # mismatch — not the "no row axis" intermediate path.
        if (
            n_idx > 0
            and log_likelihood_per_row_per_class.dim() >= 2
            and log_likelihood_per_row_per_class.shape[0] != n_idx
        ):
            raise ValueError(
                "group_index must have shape (N,) matching the "
                "leading axis of the log-likelihood; got "
                f"{flat_group_index_shape} vs N="
                f"{log_likelihood_per_row_per_class.shape[0]}"
            )
        has_row_axis = (
            n_idx > 0
            and log_likelihood_per_row_per_class.dim() >= 2
            and log_likelihood_per_row_per_class.shape[0] == n_idx
        )
    else:
        has_row_axis = False
    if not has_row_axis:
        # No N-axis: apply prior + reduce over the class axis only.
        # ``log_prior_per_group_per_class`` is broadcastable across
        # the extra axes (typically just (K,)).
        if per_group:
            raise ValueError(
                "marginalize_grouped: per_group requires a log-likelihood "
                "with a leading row axis; got shape "
                f"{tuple(log_likelihood_per_row_per_class.shape)}"
            )
        weighted = log_prior_per_group_per_class + log_likelihood_per_row_per_class
        if reduction == "logsumexp":
            return torch.logsumexp(weighted, dim=-1)
        if reduction == "sum":
            return weighted.sum(dim=-1)
        return weighted.mean(dim=-1)
    n_rows = int(log_likelihood_per_row_per_class.shape[0])
    n_classes = int(log_likelihood_per_row_per_class.shape[-1])
    extra_axes = tuple(log_likelihood_per_row_per_class.shape[1:-1])

    # Resolve single-vs-product fibration into a flat group index.
    if isinstance(group_index, tuple):
        if not isinstance(num_groups, tuple):
            raise ValueError(
                "marginalize_grouped: product fibration (tuple "
                "group_index) requires num_groups to be a tuple of "
                f"matching length; got {type(num_groups).__name__}"
            )
        flat_index, total_groups = _flatten_product_indices(group_index, num_groups)
        # Flatten an (G_1, ..., G_r, K) prior to (G_1 * ... * G_r, K).
        if log_prior_per_group_per_class.dim() == 1:
            log_prior_flat: torch.Tensor = log_prior_per_group_per_class
        elif log_prior_per_group_per_class.dim() == 2 and (
            log_prior_per_group_per_class.shape[0] == total_groups
            or log_prior_per_group_per_class.shape[0] == 1
        ):
            log_prior_flat = log_prior_per_group_per_class
        else:
            expected_axes = len(num_groups) + 1
            if log_prior_per_group_per_class.dim() != expected_axes:
                raise ValueError(
                    "marginalize_grouped: product-fibration prior "
                    f"must have shape (*{num_groups}, K) or (K,); "
                    f"got {tuple(log_prior_per_group_per_class.shape)}"
                )
            log_prior_flat = log_prior_per_group_per_class.reshape(
                total_groups, n_classes
            )
    else:
        if isinstance(num_groups, tuple):
            raise ValueError(
                "marginalize_grouped: scalar group_index requires "
                f"scalar num_groups; got {num_groups!r}"
            )
        if group_index.shape != (n_rows,):
            raise ValueError(
                "group_index must have shape (N,) matching the "
                "leading axis of the log-likelihood; got "
                f"{tuple(group_index.shape)} vs N={n_rows}"
            )
        if num_groups <= 0:
            raise ValueError(f"num_groups must be positive; got {num_groups}")
        if (group_index < 0).any() or (group_index >= num_groups).any():
            raise ValueError(
                f"group_index entries must lie in [0, {num_groups}); "
                "out-of-range index detected"
            )
        flat_index = group_index.to(torch.long)
        total_groups = int(num_groups)
        log_prior_flat = log_prior_per_group_per_class

    # Scatter-add the (N, *extra, K) per-row log-likelihood along
    # the fibration to obtain a (|G|, *extra, K) per-group
    # accumulator. ``index_add`` operates on the leading axis;
    # ``extra`` axes ride along as broadcast.
    grouped_shape: tuple[int, ...] = (total_groups,) + extra_axes + (n_classes,)
    grouped = torch.zeros(
        grouped_shape,
        dtype=log_likelihood_per_row_per_class.dtype,
        device=log_likelihood_per_row_per_class.device,
    )
    grouped = grouped.index_add(0, flat_index, log_likelihood_per_row_per_class)
    # Broadcast log_prior_flat shaped ``(K,)`` or ``(G, K)``
    # against the accumulator ``(G, *extra, K)``. We insert
    # singleton axes for ``*extra`` so the broadcast is well-formed.
    prior_view = log_prior_flat
    if prior_view.dim() == 1:
        # (K,) → (1,) * len(extra) + (K,) → broadcasts against
        # (G, *extra, K).
        for _ in extra_axes:
            prior_view = prior_view.unsqueeze(0)
    elif prior_view.dim() == 2 and prior_view.shape[0] == total_groups:
        # (G, K) → (G, *(1,)*len(extra), K).
        for _ in extra_axes:
            prior_view = prior_view.unsqueeze(1)
    weighted = prior_view + grouped
    if reduction == "logsumexp":
        per_group_aggregates = torch.logsumexp(weighted, dim=-1)
    elif reduction == "sum":
        per_group_aggregates = weighted.sum(dim=-1)
    else:  # mean
        per_group_aggregates = weighted.mean(dim=-1)
    if per_group:
        return per_group_aggregates
    # Sum over the group axis; extra axes (outer-block class
    # broadcasts) pass through unchanged.
    return per_group_aggregates.sum(dim=0)