Idea: Add Measure-Valued Pólya Urn (MVPU) Distribution for dependent sequences

[gl-alexander]

The Problem
Currently, PyMC supports discrete Bayesian nonparametrics primarily through stick-breaking representations of the Dirichlet Process. This enforces only localized reinforcement (observing category $j$ only increases the probability of $j$). When modeling various spatial, temporal, or semantic sequences, we often need non-local reinforcement, where observing category $j$ increases the probability of neighboring or correlated categories via a transition kernel.

The Proposal
I propose adding a MeasureValuedPolyaUrn sequence distribution to pymc.distributions.discrete. It generalizes the standard Pólya urn by replacing the unit-mass reinforcement with a stochastic transition kernel.

Mathematically, the predictive probability at step $n$ for $K$ discrete categories is:

$$P(X_n = j \mid X_{1:n-1}) = \frac{\mu_{0, j} + \sum_{i=1}^{n-1} R_{X_i, j}}{\sum_{k=1}^K \left( \mu_{0, k} + \sum_{i=1}^{n-1} R_{X_i, k} \right)}$$

where $\mu_0$ is the initial measure and $R$ is the $K \times K$ reinforcement matrix.

Current Prototype Status
To avoid pytensor.scan compilation bottlenecks, I've prototyped the recursive sequence by vectorizing it using pt.cumsum and indexing. NUTS can sample the continuous hyperparameters over the marginalized discrete sequence without divergences, provided the scale is anchored to ensure identifiability (e.g., using pm.Dirichlet for the initial measure and kernel proportions).

PyTensor logp logic:

# observations is a tensor of shape (N,)
# reinforcement_matrix is shape (K, K)
triggered_reinforcements = reinforcement_matrix[observations]
cum_reinforcements = pt.cumsum(triggered_reinforcements, axis=0)

# Shift cumulative sum to represent prior state at each step
shifted_cum = pt.concatenate([pt.zeros((1, K)), cum_reinforcements[:-1]], axis=0)
urn_states = initial_measure + shifted_cum

# Normalize to probabilities and extract likelihoods
predictive_probs = urn_states / pt.sum(urn_states, axis=1, keepdims=True)
seq_indices = pt.arange(observations.shape[0])
observed_probs = predictive_probs[seq_indices, observations]

Questions for the Core Team

Before I draft an official Feature Request Issue and PR with the RandomVariable Op, the check_pymcs_draws_match_reference tests, and the Sphinx docstrings:

• Does the team agree with adding this to pymc.distributions.discrete, or is a distribution of this nature better suited for pymc-experimental first?
• Are there specific PyTensor graph optimizations you prefer for sequential cumulative updates over pt.cumsum?