refactor(bass): migrate dim handling to pymc.dims

docs/source/conf.py

@@ -164,6 +164,7 @@
 # fmt: on
 numpydoc_xref_aliases = {
     "TensorVariable": ":class:`~pytensor.tensor.TensorVariable`",
+    "XTensorVariable": ":class:`~pytensor.xtensor.type.XTensorVariable`",
     "RandomVariable": ":class:`~pytensor.tensor.random.RandomVariable`",
     "ndarray": ":class:`~numpy.ndarray`",
     "DataTree": ":class:`~xarray.DataTree`",

pymc_marketing/bass/model.py

@@ -134,6 +134,7 @@
 
 """
 
+from inspect import signature
 from typing import Any, TypedDict, cast
 
 import arviz as az
@@ -142,6 +143,7 @@
 import numpy.typing as npt
 import pandas as pd
 import pymc as pm
+import pymc.dims as pmd
 import pytensor.tensor as pt
 import xarray as xr
 from matplotlib.axes import Axes
@@ -150,7 +152,13 @@
 )
 from pymc.model import Model
 from pymc.util import RandomState
-from pymc_extras.prior import Censored, Prior, VariableFactory, create_dim_handler
+from pymc_extras.prior import (
+    Censored,
+    MuAlreadyExistsError,
+    Prior,
+    VariableFactory,
+)
+from pytensor.xtensor.type import XTensorVariable
 
 from pymc_marketing.bass import plotting
 from pymc_marketing.bass.data import to_bass_dataset
@@ -159,28 +167,39 @@
 from pymc_marketing.version import __version__
 
 
+def _exp(
+    x: float | pt.TensorVariable | XTensorVariable,
+) -> pt.TensorVariable | XTensorVariable:
+    """``exp`` that works for floats, PyTensor tensors, and xtensor variables.
+
+    Lets :func:`F` and :func:`f` be used both with NumPy/float inputs and with
+    the named-dims (xtensor) variables of the ``pymc.dims`` model graph.
+    """
+    return pmd.math.exp(x) if isinstance(x, XTensorVariable) else pt.exp(x)
+
+
 def F(
-    p: float | pt.TensorVariable,
-    q: float | pt.TensorVariable,
-    t: float | pt.TensorVariable,
-) -> pt.TensorVariable:
+    p: float | pt.TensorVariable | XTensorVariable,
+    q: float | pt.TensorVariable | XTensorVariable,
+    t: float | pt.TensorVariable | XTensorVariable,
+) -> pt.TensorVariable | XTensorVariable:
     r"""Installed base fraction (cumulative adoption proportion).
 
     This function calculates the cumulative proportion of adopters at time t,
     representing the fraction of the potential market that has adopted the product.
 
     Parameters
     ----------
-    p : float or TensorVariable
+    p : float, TensorVariable or XTensorVariable
         Coefficient of innovation (external influence)
-    q : float or TensorVariable
+    q : float, TensorVariable or XTensorVariable
         Coefficient of imitation (internal influence)
-    t : array-like or TensorVariable
+    t : array-like, TensorVariable or XTensorVariable
         Time points
 
     Returns
     -------
-    TensorVariable
+    TensorVariable or XTensorVariable
         The cumulative proportion of adopters at each time point
 
     Notes
@@ -193,14 +212,14 @@ def F(
 
     When :math:`t=0`, :math:`F(t)=0`, and as :math:`t` approaches infinity, :math:`F(t)` approaches 1.
     """
-    return (1 - pt.exp(-(p + q) * t)) / (1 + (q / p) * pt.exp(-(p + q) * t))
+    return (1 - _exp(-(p + q) * t)) / (1 + (q / p) * _exp(-(p + q) * t))
 
 
 def f(
-    p: float | pt.TensorVariable,
-    q: float | pt.TensorVariable,
-    t: float | pt.TensorVariable,
-) -> pt.TensorVariable:
+    p: float | pt.TensorVariable | XTensorVariable,
+    q: float | pt.TensorVariable | XTensorVariable,
+    t: float | pt.TensorVariable | XTensorVariable,
+) -> pt.TensorVariable | XTensorVariable:
     r"""Installed base fraction rate of change (adoption rate).
 
     This function calculates the rate of new adoptions at time t as a
@@ -209,16 +228,16 @@ def f(
 
     Parameters
     ----------
-    p : float or TensorVariable
+    p : float, TensorVariable or XTensorVariable
         Coefficient of innovation (external influence)
-    q : float or TensorVariable
+    q : float, TensorVariable or XTensorVariable
         Coefficient of imitation (internal influence)
-    t : array-like or TensorVariable
+    t : array-like, TensorVariable or XTensorVariable
         Time points
 
     Returns
     -------
-    TensorVariable
+    TensorVariable or XTensorVariable
         The adoption rate at each time point as a fraction of potential market
 
     Notes
@@ -237,9 +256,103 @@ def f(
 
     The peak adoption rate occurs at time :math:`t^* = \frac{\ln(q/p)}{p+q}`
     """
-    return (p * pt.square(p + q) * pt.exp(t * (p + q))) / pt.square(
-        p * pt.exp(t * (p + q)) + q
+    exp_t = _exp(t * (p + q))
+    return (p * (p + q) ** 2 * exp_t) / (p * exp_t + q) ** 2
+
+
+def _supports_xdist(prior: Prior | Censored | VariableFactory) -> bool:
+    """Whether ``prior`` can build itself directly in ``pymc.dims`` space."""
+    if "xdist" not in signature(prior.create_variable).parameters:
+        # A VariableFactory written against the ``create_variable(name)``
+        # signature the protocol documents.
+        return False
+
+    if isinstance(prior, Censored):
+        return hasattr(pmd, "Censored") and _supports_xdist(prior.distribution)
+
+    if isinstance(prior, Prior):
+        return hasattr(pmd, prior.distribution)
+
+    # A factory that takes the kwarg handles its own dispatch.
+    return True
+
+
+def _create_dim_variable(
+    prior: Prior | Censored | VariableFactory, name: str
+) -> XTensorVariable:
+    """Create a named-dims (xtensor) variable for a prior.
+
+    Uses ``xdist=True`` so the variable lives in xtensor space from the start.
+    Whatever cannot take that route -- a distribution ``pymc.dims`` does not
+    implement (``DiracDelta`` before pymc 6.2, ``Wald``, ...) or a custom
+    factory written against the ``create_variable(name)`` signature -- gets a
+    regular variable wrapped with ``pmd.as_xtensor``.
+    """
+    if not _supports_xdist(prior):
+        return pmd.as_xtensor(prior.create_variable(name), dims=prior.dims or ())
+    return prior.create_variable(name, xdist=True)
+
+
+def _create_likelihood_variable(
+    prior: Prior | Censored,
+    name: str,
+    mu: XTensorVariable,
+    observed: XTensorVariable | None,
+    dims: tuple[str, ...],
+) -> pt.TensorVariable | XTensorVariable:
+    """Create the outcome variable, observed or not.
+
+    A distribution ``pymc.dims`` does not implement gets the regular pymc
+    path, on the underlying tensors in ``dims`` order.
+
+    ``create_likelihood_variable`` is for the observed case only: a
+    likelihood needs data, so pymc_extras refuses ``observed=None`` there
+    (pymc-devs/pymc-extras#731). Prior predictive still needs the outcome
+    node, so build it with ``create_variable`` and ``mu`` attached, keeping
+    the same ``mu`` guard the pymc_extras method applies.
+    """
+    xdist = _supports_xdist(prior)
+
+    if observed is not None:
+        if xdist:
+            return prior.create_likelihood_variable(
+                name, mu=mu, observed=observed, xdist=True
+            )
+        return prior.create_likelihood_variable(
+            name,
+            mu=mu.transpose(*dims).values,
+            observed=observed.transpose(*dims).values,
+        )
+
+    # Censored keeps its parameters on the wrapped distribution.
+    inner = prior.distribution if isinstance(prior, Censored) else prior
+    if "mu" in inner.parameters:
+        raise MuAlreadyExistsError(inner)
+
+    unobserved = inner.deepcopy()
+    unobserved.parameters["mu"] = mu if xdist else mu.transpose(*dims).values
+    outcome: Prior | Censored = (
+        Censored(unobserved, lower=prior.lower, upper=prior.upper)
+        if isinstance(prior, Censored)
+        else unobserved
     )
+    return outcome.create_variable(name, xdist=xdist)
+
+
+def _align_to_dims(
+    var: XTensorVariable, dims: tuple[str, ...], model: Model
+) -> XTensorVariable:
+    """Broadcast ``var`` up to ``dims``, in that order.
+
+    ``pmd.Deterministic`` takes the dims off the graph and only transposes
+    them, so a quantity built from parameters that do not carry every model
+    dim would be stored without those dims. Sizes come from the model, which
+    knows every dim it has registered, coords argument or not.
+    """
+    missing = {dim: model.dim_lengths[dim] for dim in dims if dim not in var.dims}
+    if missing:
+        var = var.expand_dims(missing)
+    return var.transpose(*dims)
 
 
 class BassPriors(TypedDict):
@@ -319,47 +432,59 @@ def create_bass_model(
     """
     model = model or pm.Model(coords=coords)
     with model:
-        parameter_dims = (
-            set(priors["p"].dims or ())
-            .union(priors["q"].dims or ())
-            .union(priors["m"].dims or ())
+        # Declaration order, not set order: `combined_dims` labels the axes of
+        # `observed` positionally, so an order that varies between processes
+        # would silently mislabel the data.
+        declared_dims = (
+            *(priors["p"].dims or ()),
+            *(priors["q"].dims or ()),
+            *(priors["m"].dims or ()),
+            *(getattr(priors["likelihood"], "dims", ()) or ()),
         )
-        likelihood_dims = set(getattr(priors["likelihood"], "dims", ()) or ())
-
         combined_dims = (
             "T",
-            *tuple(parameter_dims.union(likelihood_dims).difference(["T"])),
+            *(dim for dim in dict.fromkeys(declared_dims) if dim != "T"),
         )
-        dim_handler = create_dim_handler(combined_dims)
 
-        m = dim_handler(priors["m"].create_variable("m"), priors["m"].dims)
-        p = dim_handler(priors["p"].create_variable("p"), priors["p"].dims)
-        q = dim_handler(priors["q"].create_variable("q"), priors["q"].dims)
+        time = pmd.as_xtensor(t, dims=("T",))
+        m = _create_dim_variable(priors["m"], "m")
+        p = _create_dim_variable(priors["p"], "p")
+        q = _create_dim_variable(priors["q"], "q")
 
-        time = dim_handler(t, "T")
+        def deterministic(name: str, value: XTensorVariable) -> XTensorVariable:
+            return pmd.Deterministic(name, _align_to_dims(value, combined_dims, model))
 
-        adopters = pm.Deterministic("adopters", m * f(p, q, time), dims=combined_dims)
+        adopters = deterministic("adopters", m * f(p, q, time))
+        deterministic("innovators", m * p * (1 - F(p, q, time)))
+        deterministic("imitators", m * q * F(p, q, time) * (1 - F(p, q, time)))
 
-        pm.Deterministic(
-            "innovators",
-            m * p * (1 - F(p, q, time)),
-            dims=combined_dims,
-        )
-        pm.Deterministic(
-            "imitators",
-            m * q * F(p, q, time) * (1 - F(p, q, time)),
-            dims=combined_dims,
+        # `peak` carries only the parameter dims, but in `combined_dims` order
+        # so it agrees with the deterministics above.
+        peak = (pmd.math.log(q) - pmd.math.log(p)) / (p + q)
+        pmd.Deterministic(
+            "peak", peak.transpose(*(dim for dim in combined_dims if dim in peak.dims))
         )
 
-        peak = (pt.log(q) - pt.log(p)) / (p + q)
-        peak_dims = tuple(parameter_dims) if parameter_dims else None
-        pm.Deterministic("peak", peak, dims=peak_dims)
-
         priors["likelihood"].dims = combined_dims
-        priors["likelihood"].create_likelihood_variable(  # type: ignore
+        if observed is None:
+            observed_xt = None
+        else:
+            # The data knows its own axis labels, which `combined_dims` need
+            # not match: an `xr.DataArray` carries them, a registered
+            # `pm.Data` has them on the model.
+            observed_dims = getattr(observed, "dims", None) or (
+                model.named_vars_to_dims.get(
+                    getattr(observed, "name", None), combined_dims
+                )
+            )
+            observed_xt = pmd.as_xtensor(observed, dims=tuple(observed_dims))
+
+        _create_likelihood_variable(
+            priors["likelihood"],
             "y",
             mu=adopters,
-            observed=observed,
+            observed=observed_xt,
+            dims=combined_dims,
         )
 
     return model

pymc_marketing/mlflow.py

@@ -562,7 +562,11 @@ def log_model_graph(model: Model, path: str | Path) -> None:
 
 def _get_random_variable_name(rv) -> str:
     # Taken from new version of pymc/model_graph.py
-    symbol = rv.owner.op.__class__.__name__
+    op = rv.owner.op
+    # A `pymc.dims` variable wraps the real RV in a generic `XRV`, which would
+    # otherwise be reported as "X". The distribution is on the wrapped op.
+    op = getattr(op, "core_op", op)
+    symbol = op.__class__.__name__
 
     if symbol.endswith("RV"):
         symbol = symbol[:-2]

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "pandas",
     "patsy",
     "pydantic>=2.1.0",
-    "pymc>=6.0.0,<6.1.0",
+    "pymc>=6.0.1,<6.1.0",
     "pytensor>=3.0.0,<4.0",
     "scikit-learn>=1.1.1",
     "seaborn>=0.12.2",