From b22c382b36679312ef27ab1b91cbab1b714e6c09 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Fri, 26 Jun 2026 22:51:52 +0200 Subject: [PATCH] Recognize affine functions of independent Gaussians as MvNormal logp Teach pymc.logprob to derive the logp of a value variable that is an affine function of one or more independent Gaussian leaves (e.g. the low-rank guide loc + W @ ek + d * ed) by propagating moments and substituting a dense MvNormal, instead of inverting the transform. This unlocks fundamentally non-invertible cases the existing machinery declines: wide/low-rank matvecs (find_measurable_matmul bails on non-square matrices) and sums of >=2 independent Gaussians (the elemwise transforms require <=1 measurable operand). Covers low-rank ADVI guides, sums of independent normals, and linear-Gaussian marginalization (regression coefficients, etc.). The accumulated covariance is kept dense; exploiting low-rank structure for an O(D K^2) logp would need a Woodbury rewrite in PyTensor, which does not exist yet, so the emitted MvNormal does a dense O(D^3) cholesky. Claude-Session: https://claude.ai/code/session_0183HnSTpYk7UrvAJ7D3kBtj --- pymc/logprob/__init__.py | 1 + pymc/logprob/gaussian.py | 313 +++++++++++++++++++++++++++++++++ tests/logprob/test_gaussian.py | 169 ++++++++++++++++++ 3 files changed, 483 insertions(+) create mode 100644 pymc/logprob/gaussian.py create mode 100644 tests/logprob/test_gaussian.py diff --git a/pymc/logprob/__init__.py b/pymc/logprob/__init__.py index 985f9f489d..c62cdd0b94 100644 --- a/pymc/logprob/__init__.py +++ b/pymc/logprob/__init__.py @@ -51,6 +51,7 @@ import pymc.logprob.arithmetic import pymc.logprob.cumsum import pymc.logprob.checks +import pymc.logprob.gaussian import pymc.logprob.linalg import pymc.logprob.mixture import pymc.logprob.order diff --git a/pymc/logprob/gaussian.py b/pymc/logprob/gaussian.py new file mode 100644 index 0000000000..c4abc3dd31 --- /dev/null +++ b/pymc/logprob/gaussian.py @@ -0,0 +1,313 @@ +# Copyright 2024 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Recognize a value variable that is an affine function of independent Gaussian leaves. + +A linear function of independent Gaussians is itself Gaussian, so its density is +available in closed form even when the map is non-invertible (wide / low-rank +matrices, sums of independent normals) and the usual invert-and-Jacobian +machinery (``MeasurableMatMul``, the elemwise transforms) cannot derive it. + +Instead of inverting, we *propagate moments*: walk the affine subgraph feeding a +value variable, accumulate a symbolic ``(mean, cov)`` over its last axis, and +substitute a plain ``MvNormal`` whose existing ``_logprob`` produces the density. + +The accumulated covariance is kept dense (e.g. ``W @ W.T + diag(d**2)`` for the +low-rank case). Exploiting its structure for an ``O(D K**2)`` logp would require a +Woodbury / matrix-determinant-lemma rewrite in PyTensor, which does not exist +yet; until then the emitted ``MvNormal`` logp does a dense ``O(D**3)`` cholesky. +""" + +import pytensor.tensor as pt + +from pytensor.graph.rewriting.basic import node_rewriter +from pytensor.graph.traversal import ancestors +from pytensor.scalar.basic import Add as ScalarAdd +from pytensor.scalar.basic import Mul as ScalarMul +from pytensor.tensor.elemwise import DimShuffle, Elemwise +from pytensor.tensor.math import _matmul +from pytensor.tensor.random.basic import MvNormalRV, NormalRV, multivariate_normal +from pytensor.tensor.random.op import RandomVariable +from pytensor.tensor.variable import TensorVariable + +from pymc.logprob.abstract import MeasurableOp +from pymc.logprob.rewriting import measurable_ir_rewrites_db +from pymc.logprob.transforms import LocTransform, MeasurableTransform, ScaleTransform +from pymc.logprob.utils import check_potential_measurability + +# A moment triple ``(mean, cov, leaves)`` describes a variable as Gaussian over its +# last axis: ``mean`` has shape ``(..., n)``, ``cov`` shape ``(..., n, n)``, and +# ``leaves`` is the frozenset of RandomVariable outputs feeding it (used to check +# that independent terms depend on disjoint leaves so their covariances add). + + +def _diag_from_vec(v): + """Build a ``(..., n, n)`` matrix with diagonal ``v`` (shape ``(..., n)``).""" + return pt.eye(v.shape[-1]) * v[..., None, :] + + +def _rv_leaves(var) -> frozenset: + return frozenset( + a + for a in ancestors([var]) + if a.owner is not None and isinstance(a.owner.op, RandomVariable) + ) + + +def _is_squeeze_last_axis(op, inp) -> bool: + """True if ``op`` is a DimShuffle that just drops the last (length-1) axis.""" + if not isinstance(op, DimShuffle): + return False + return op.new_order == tuple(range(inp.type.ndim - 1)) + + +def _is_expand_last_axis(op) -> bool: + """True if ``op`` is a DimShuffle that just appends a new last axis.""" + if not isinstance(op, DimShuffle): + return False + new_order = op.new_order + return ( + len(new_order) == op.input_ndim + 1 + and new_order[-1] == "x" + and new_order[:-1] == tuple(range(op.input_ndim)) + ) + + +def _normal_leaf_moments(var): + """Moments of an independent ``Normal`` vector over its last axis.""" + if var.type.ndim < 1: + return None + rng, size, mu, sigma = var.owner.inputs + shape = var.shape + mean = pt.broadcast_to(mu, shape) + cov = _diag_from_vec(pt.broadcast_to(sigma, shape) ** 2) + return mean, cov, frozenset({var}) + + +def _column_moments(col): + """Moments of a column Gaussian ``(..., n, 1)`` along its ``n`` axis. + + Returns ``(mean_col, cov, leaves)`` with ``mean_col`` shape ``(..., n, 1)`` and + ``cov`` shape ``(..., n, n)``, or ``None``. + """ + node = col.owner + if node is None: + return None + op = node.op + + # ``ExpandDims(g)`` of a vector Gaussian (un-lifted form). + if _is_expand_last_axis(op): + inner = _affine_gaussian_moments(node.inputs[0]) + if inner is None: + return None + mean_v, cov, leaves = inner + return mean_v[..., None], cov, leaves + + # A ``Normal`` RV lifted directly into ``(..., n, 1)`` shape. + if isinstance(op, NormalRV): + rng, size, mu, sigma = node.inputs + shape = col.shape + mean_col = pt.broadcast_to(mu, shape) + cov = _diag_from_vec((pt.broadcast_to(sigma, shape) ** 2)[..., 0]) + return mean_col, cov, frozenset({col}) + + return None + + +def _moments_of_matvec(var): + """Moments of ``A @ g`` written as ``Squeeze(Matmul(A, g_col))``.""" + node = var.owner + if not _is_squeeze_last_axis(node.op, node.inputs[0]): + return None + mm = node.inputs[0] + if mm.owner is None or mm.owner.op != _matmul: + return None + A, col = mm.owner.inputs + # The constant operand must carry no hidden measurable dependency. + if check_potential_measurability([A]): + return None + inner = _column_moments(col) + if inner is None: + return None + mean_col, cov_g, leaves = inner + mean = (A @ mean_col)[..., 0] + cov = A @ cov_g @ A.mT + return mean, cov, leaves + + +def _moments_of_add(node, *, require_two=False): + """Moments of a sum of independent Gaussian terms plus constant shifts.""" + shifts = [] + terms = [] + for inp in node.inputs: + if not _rv_leaves(inp): + shifts.append(inp) + continue + moments = _affine_gaussian_moments(inp) + if moments is None: + return None + terms.append(moments) + + if not terms or (require_two and len(terms) < 2): + return None + + # Covariances only add if the terms depend on disjoint leaves. + all_leaves: frozenset = frozenset() + for _, _, leaves in terms: + if all_leaves & leaves: + return None + all_leaves |= leaves + + mean = terms[0][0] + cov = terms[0][1] + for term_mean, term_cov, _ in terms[1:]: + mean = mean + term_mean + cov = cov + term_cov + for shift in shifts: + mean = mean + shift + return mean, cov, all_leaves + + +def _moments_of_scale(node): + """Moments of an elementwise ``c * g`` with constant ``c``.""" + gaussian = None + consts = [] + for inp in node.inputs: + if _rv_leaves(inp): + if gaussian is not None: + return None + gaussian = inp + else: + consts.append(inp) + if gaussian is None: + return None + inner = _affine_gaussian_moments(gaussian) + if inner is None: + return None + mean, cov, leaves = inner + c = consts[0] if len(consts) == 1 else pt.mul(*consts) + mean = c * mean + cov = cov * c[..., None, :] * c[..., :, None] + return mean, cov, leaves + + +def _moments_through_transform(node): + """Read moments through an already-measurable Loc/Scale transform.""" + op = node.op + transform = op.transform_elemwise + base = node.inputs[op.measurable_input_idx] + inner = _affine_gaussian_moments(base) + if inner is None: + return None + mean, cov, leaves = inner + others = [inp for i, inp in enumerate(node.inputs) if i != op.measurable_input_idx] + if isinstance(transform, LocTransform): + loc = others[0] if len(others) == 1 else pt.add(*others) + return mean + loc, cov, leaves + if isinstance(transform, ScaleTransform): + c = others[0] if len(others) == 1 else pt.mul(*others) + return c * mean, cov * c[..., None, :] * c[..., :, None], leaves + # Any other transform (Exp, Log, Abs, ...) is non-affine. + return None + + +def _affine_gaussian_moments(var): + """Interpret ``var`` as an affine function of independent Gaussian leaves. + + Returns the dense moment triple ``(mean, cov, leaves)`` over ``var``'s last + axis, or ``None`` if ``var`` is not such an affine function. + """ + node = var.owner + if node is None: + return None + op = node.op + + if isinstance(op, NormalRV): + return _normal_leaf_moments(var) + + if isinstance(op, MvNormalRV): + rng, size, mu, cov = node.inputs + return pt.broadcast_arrays(mu, cov[..., -1])[0], cov, frozenset({var}) + + if isinstance(op, MeasurableTransform): + return _moments_through_transform(node) + + if isinstance(op, DimShuffle): + return _moments_of_matvec(var) + + if isinstance(op, Elemwise): + if isinstance(op.scalar_op, ScalarAdd): + return _moments_of_add(node) + if isinstance(op.scalar_op, ScalarMul): + return _moments_of_scale(node) + + return None + + +def _square_or_unknown(A) -> bool: + """True unless ``A`` is statically known to be non-square (so we should fire).""" + m, n = A.type.shape[-2], A.type.shape[-1] + return m is None or n is None or m == n + + +def _emit_mvnormal(mean, cov, leaves) -> list[TensorVariable]: + rng = next((leaf.owner.inputs[0] for leaf in leaves), None) + return [multivariate_normal(mean, cov, rng=rng)] + + +@node_rewriter(tracks=[DimShuffle]) +def find_measurable_matvec_normal(fgraph, node): + """Recognize ``A @ g`` of a Gaussian with non-square ``A`` as an ``MvNormal``.""" + if isinstance(node.op, MeasurableOp): + return None + if not _is_squeeze_last_axis(node.op, node.inputs[0]): + return None + mm = node.inputs[0] + if mm.owner is None or mm.owner.op != _matmul: + return None + A, _ = mm.owner.inputs + # Square (or statically unknown) maps are invertible; leave them to + # MeasurableMatMul to avoid stepping on that path. + if _square_or_unknown(A): + return None + moments = _affine_gaussian_moments(node.outputs[0]) + if moments is None: + return None + return _emit_mvnormal(*moments) + + +@node_rewriter(tracks=[Elemwise]) +def find_measurable_sum_of_gaussians(fgraph, node): + """Recognize a sum of >=2 independent Gaussians as an ``MvNormal``.""" + if isinstance(node.op, MeasurableOp): + return None + if not isinstance(node.op.scalar_op, ScalarAdd): + return None + moments = _moments_of_add(node, require_two=True) + if moments is None: + return None + return _emit_mvnormal(*moments) + + +measurable_ir_rewrites_db.register( + find_measurable_matvec_normal.__name__, + find_measurable_matvec_normal, + "basic", + "gaussian", +) +measurable_ir_rewrites_db.register( + find_measurable_sum_of_gaussians.__name__, + find_measurable_sum_of_gaussians, + "basic", + "gaussian", +) diff --git a/tests/logprob/test_gaussian.py b/tests/logprob/test_gaussian.py new file mode 100644 index 0000000000..fb4ee6ade4 --- /dev/null +++ b/tests/logprob/test_gaussian.py @@ -0,0 +1,169 @@ +# Copyright 2024 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytensor.tensor as pt +import pytest + +from pymc.distributions import Exponential, MvNormal, Normal +from pymc.logprob.basic import conditional_logp, logp + + +def assert_logp_equal(y, ref, test_value, subs=None): + """Total logp of ``y`` matches the reference distribution at ``test_value``.""" + subs = subs or {} + value = y.type() + y_logp = logp(y, value) + ref_logp = logp(ref, value) + np.testing.assert_allclose( + y_logp.eval({value: test_value, **subs}).sum(), + ref_logp.eval({value: test_value, **subs}).sum(), + ) + + +@pytest.mark.parametrize("D, K", [(4, 2), (5, 1), (3, 3)]) +def test_low_rank_guide(D, K): + """The low-rank ADVI guide ``loc + W @ ek + d * ed`` is recognized as MvNormal.""" + rng = np.random.default_rng(623) + loc_v = rng.normal(size=D) + W_v = rng.normal(size=(D, K)) + d_v = np.abs(rng.normal(size=D)) + 0.5 + + loc = pt.tensor("loc", shape=(D,)) + W = pt.tensor("W", shape=(D, K)) + d = pt.tensor("d", shape=(D,)) + ek = Normal.dist(0, 1, shape=(K,)) + ed = Normal.dist(0, 1, shape=(D,)) + u = loc + W @ ek + d * ed + + ref = MvNormal.dist(mu=loc, cov=W @ W.T + pt.diag(d**2)) + assert_logp_equal(u, ref, rng.normal(size=D), {loc: loc_v, W: W_v, d: d_v}) + + +def test_sum_of_two_full_normals(): + """``loc + A @ z1 + B @ z2`` -> MvNormal(loc, A Aᵀ + B Bᵀ).""" + rng = np.random.default_rng(1) + D = 5 + loc_v, A_v, B_v = rng.normal(size=D), rng.normal(size=(D, D)), rng.normal(size=(D, D)) + + loc = pt.tensor("loc", shape=(D,)) + A = pt.tensor("A", shape=(D, D)) + B = pt.tensor("B", shape=(D, D)) + z1 = Normal.dist(0, 1, shape=(D,)) + z2 = Normal.dist(0, 1, shape=(D,)) + u = loc + A @ z1 + B @ z2 + + ref = MvNormal.dist(mu=loc, cov=A @ A.T + B @ B.T) + assert_logp_equal(u, ref, rng.normal(size=D), {loc: loc_v, A: A_v, B: B_v}) + + +def test_sum_of_two_diagonal_normals(): + """Adding two independent diagonal Normals yields the joint (diagonal) MvNormal.""" + rng = np.random.default_rng(2) + D = 4 + x = Normal.dist(mu=1.0, sigma=2.0, shape=(D,)) + y = Normal.dist(mu=-0.5, sigma=0.7, shape=(D,)) + ref = MvNormal.dist(mu=np.full(D, 0.5), cov=np.eye(D) * (2.0**2 + 0.7**2)) + assert_logp_equal(x + y, ref, rng.normal(size=D)) + + +def test_mvnormal_plus_normal_promotion(): + """A full MvNormal plus an independent diagonal Normal stays MvNormal.""" + rng = np.random.default_rng(3) + D = 4 + A_v = rng.normal(size=(D, D)) + cov_v = A_v @ A_v.T + np.eye(D) + mv = MvNormal.dist(mu=np.zeros(D), cov=cov_v) + ed = Normal.dist(0, 0.5, shape=(D,)) + ref = MvNormal.dist(mu=np.zeros(D), cov=cov_v + np.eye(D) * 0.25) + assert_logp_equal(mv + ed, ref, rng.normal(size=D)) + + +def test_linear_regression_marginal(): + """Analytic marginal ``y = X @ beta + sigma * eps`` over Gaussian latents.""" + rng = np.random.default_rng(4) + N, P, sigma = 6, 3, 0.8 + X_v = rng.normal(size=(N, P)) + mu_beta = rng.normal(size=P) + sd_beta = np.abs(rng.normal(size=P)) + 0.3 + + X = pt.tensor("X", shape=(N, P)) + beta = Normal.dist(mu=pt.as_tensor(mu_beta), sigma=pt.as_tensor(sd_beta), shape=(P,)) + eps = Normal.dist(0, 1, shape=(N,)) + y = X @ beta + sigma * eps + + ref_cov = X @ pt.as_tensor(np.diag(sd_beta**2)) @ X.T + sigma**2 * pt.eye(N) + ref = MvNormal.dist(mu=X @ pt.as_tensor(mu_beta), cov=ref_cov) + assert_logp_equal(y, ref, rng.normal(size=N), {X: X_v}) + + +def test_batched_low_rank(): + """Leading batch dims are carried through the moment propagation.""" + rng = np.random.default_rng(5) + Bz, D, K = 4, 5, 2 + loc_v = rng.normal(size=(Bz, D)) + W_v = rng.normal(size=(Bz, D, K)) + d_v = np.abs(rng.normal(size=(Bz, D))) + 0.5 + + loc = pt.tensor("loc", shape=(Bz, D)) + W = pt.tensor("W", shape=(Bz, D, K)) + d = pt.tensor("d", shape=(Bz, D)) + ek = Normal.dist(0, 1, shape=(Bz, K)) + ed = Normal.dist(0, 1, shape=(Bz, D)) + u = loc + (W @ ek[..., None])[..., 0] + d * ed + + ref = MvNormal.dist(mu=loc, cov=W @ W.mT + pt.eye(D) * d[..., None, :] ** 2) + assert_logp_equal(u, ref, rng.normal(size=(Bz, D)), {loc: loc_v, W: W_v, d: d_v}) + + +def test_full_rank_square_still_uses_matmul(): + """A square ``L @ z`` keeps deriving (via MeasurableMatMul), unaffected.""" + rng = np.random.default_rng(6) + D = 3 + loc_v, L_v = rng.normal(size=D), rng.normal(size=(D, D)) + + loc = pt.tensor("loc", shape=(D,)) + L = pt.tensor("L", shape=(D, D)) + z = Normal.dist(0, 1, shape=(D,)) + u = loc + L @ z + + ref = MvNormal.dist(mu=loc, cov=L @ L.T) + assert_logp_equal(u, ref, rng.normal(size=D), {loc: loc_v, L: L_v}) + + +@pytest.mark.parametrize( + "build", + [ + # non-Gaussian leaf in the affine combination + lambda D, K, loc, W, d, ed: loc + W @ Exponential.dist(1.0, shape=(K,)) + d * ed, + # non-linear op inside the affine path + lambda D, K, loc, W, d, ed: loc + W @ Normal.dist(0, 1, shape=(K,)) + pt.exp(ed), + ], +) +def test_bails_cleanly(build): + D, K = 5, 2 + loc = pt.tensor("loc", shape=(D,)) + W = pt.tensor("W", shape=(D, K)) + d = pt.tensor("d", shape=(D,)) + ed = Normal.dist(0, 1, shape=(D,)) + y = build(D, K, loc, W, d, ed) + with pytest.raises((NotImplementedError, RuntimeError)): + conditional_logp({y: y.type()}) + + +def test_bails_on_correlated_leaves(): + """Sharing a leaf across summands violates independence -> bail.""" + D = 4 + z = Normal.dist(0, 1, shape=(D,)) + with pytest.raises((NotImplementedError, RuntimeError)): + conditional_logp({z + 2 * z: pt.vector("v", shape=(D,))})