Bayesian Rank-1 Non-Negative Canonical Polyadic Decomposition In PyMC

This post is just a quick example of how to perform Bayesian non-negative canonical polyadic decomposition in PyMC. Here we will only perform a rank 1 decomposition with unit coefficients.

Let us suppose that we have a sequence of random vectors

\[\vec u_1, \ldots, \vec u_k\]

\[\vec u_i \sim \text{Exponential}(\vec \lambda_i)\]

then their outer product makes

\[\mathcal{L} = \bigotimes_{i=1}^k \vec u_i\]

and we’ll suppose that

\[\mathcal{X} \sim \text{Exponential}(\mathcal{L})\]

is our observed random variable.

"""Model construction for toy Bayesian CPD."""

from __future__ import annotations

import numpy as np
import pymc as pm
from pytensor import tensor as at


def simulate_data(size: int = 3) -> tuple[np.ndarray, int]:
    """Simulate prior-predictive data for the toy example."""
    rng = np.random.default_rng(0)
    u1 = rng.exponential(0.001, size=size)
    u2 = rng.exponential(0.001, size=size)
    u3 = rng.exponential(0.001, size=size)
    x = np.outer(u1, u2)
    x = np.outer(x, u3).reshape((size,) * 3).flatten()
    return x, size


def build_model(size: int = 3, observed: np.ndarray | None = None) -> pm.Model:
    """Build and return the PyMC model."""
    model = pm.Model()
    if observed is None:
        observed, _ = simulate_data(size)

    with model:
        u1 = pm.Exponential("u1", lam=1, size=size)
        u2 = pm.Exponential("u2", lam=1, size=size)
        u3 = pm.Exponential("u3", lam=1, size=size)
        mu = at.outer(at.outer(u1, u2), u3).flatten()

        pm.Exponential("X", mu, shape=observed.size, observed=observed)
    return model

The more general cases of finite rank-\(r\) decomposition should be built with a function or class that constructs it. One option is pymc_experimental.model_builder.ModelBuilder. I imagine such a class should have a parameter for the rank, a parameter for distribution(S) for the random vectors, and also a parameter for the distribution of the observations. A complicated case would be to specify the distribution of every random variable, but then the config alone for the model would be complicated. A possible partial solution for that is to default to assume that the input specifies a distribution family for all the random vectors, or if a dict is passed specifying modes then those distributions are used over the mode. I think specifying the distribution of every coefficient parameter and every element of every vector is sufficiently “custom” that one might as well use the default PyMC model construction class in a context wrapper.