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

[anevolbap]

Closes #2598

Migrates the Bass model's dimension handling from create_dim_handler (pymc_extras.prior) to pymc.dims, for consistency with the rest of the codebase.

Reopens #2724, which GitHub auto-closed when the v1.0.0 branch was deleted and merged into main. Same work, retargeted to main.

In create_bass_model:

• m, p, q and the likelihood via create_variable(xdist=True); deterministics via pmd.Deterministic.
• The public F and f take xtensor inputs, so there is no second copy of the formulas for the model graph.
• t and y_obs stay pm.Data wrapped with pmd.as_xtensor, so the data-setter and out-of-sample forecasting keep working.
• Deterministics carry the dims their parameters have, transposed to (T, ...). They are not broadcast up to the likelihood's dims: pooled parameters give pooled curves, and a per-product curve is the user's to build outside the model.
• This is a user-visible change: on main the deterministics carried the union of the p, q and m dims.
• Without data (prior predictive), the outcome variable is built with create_variable and mu set, keeping the UnsupportedDistributionError and MuAlreadyExistsError guards, since upstream refuses observed=None (Prior.create_likelihood_variable(observed=None) fails on the xdist path pymc-devs/pymc-extras#731).
• Priors must use distributions pymc.dims implements, and a custom VariableFactory must take the xdist keyword.

DiracDelta is native in pymc.dims on the pymc>=6.2 floor, so #2726 has nothing left to remove and can be closed.

Verified: full bass test suite, mlflow autolog, adopters/peak matching the closed-form formula to machine precision, and the bass notebook mock runner.

---

📚 Documentation preview 📚: https://pymc-marketing--2771.org.readthedocs.build/en/2771/

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.83%. Comparing base (66d575d) to head (18a2e78).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2771      +/-   ##
==========================================
+ Coverage   94.81%   94.83%   +0.01%     
==========================================
  Files         111      111              
  Lines       17623    17665      +42     
==========================================
+ Hits        16710    16752      +42     
  Misses        913      913              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[williambdean]

Updating the base branch. Should work now with pymc version

[williambdean]

Looking at the CI, the oldest-deps shard installs pymc==6.0.0 which lacks pmd.Poisson and pmd.NegativeBinomial. Bump the floor to 6.0.1 in pyproject.toml and it should be green.

[anevolbap]

Addressed in 1bc024c:

1. Added a bullet to the PR description flagging the peak shape change (dims narrowed to what p and q declare, m-only indexing stops working).
2. Added test_observed_data_without_dims_is_labelled_positionally, checking that a pm.Data without dims gets the same labelling and logp as one registered with explicit dims, plus a note in the observed docstring about the positional fallback.
3. Added a sentence to the _supports_xdist docstring: a factory that hides xdist behind **kwargs takes the wrapping path, declaring the parameter explicitly takes the native one.

9f7af31 merges main to resolve the pyproject.toml/uv.lock conflict from #2883; the dependency files now match main and dropped out of the diff.

[daimon-pymclabs]

Checked 1bc024c and 9f7af31. All three notes are closed.

The PR description now flags the peak dim narrowing and the loss of m-only indexing, which was the only item anyone downstream needed warning about. test_observed_data_without_dims_is_labelled_positionally pins the fallback the way I wanted it pinned, comparing labelling and logp against the explicitly-dimmed pm.Data, and the observed docstring now states where labels are read from and what the positional (T, ...) order means. The _supports_xdist docstring says the check reads the literal signature, so **kwargs gets the wrapping path.

The merge drops pyproject.toml and uv.lock out of the diff and CI is green apart from the two skipped PyPI upload jobs. Nothing further from me.

[anevolbap]

create_likelihood_variable is only called with data: pymc-extras refuses observed=None there since pymc-devs/pymc-extras#732. The unobserved node goes through create_variable with mu attached instead, so the two guards that method applies (UnsupportedDistributionError for a mu-less distribution, MuAlreadyExistsError) are mirrored here to keep both paths raising the same typed errors. Covered by test_mu_less_likelihood_without_observed_raises and test_unobserved_skips_create_likelihood_variable.

[anevolbap]

Removed the # Taken from new version of pymc/model_graph.py line here: it dated to a900f97 (2024) and pymc's random_variable_symbol has since diverged (it now prefers op.name with a class-name fallback, and has no core_op unwrap), so the reference no longer described this helper. The core_op line is local to this PR.

[williambdean]

Thanks for all of the iteraetion on this, @anevolbap! I think we are over complicating a bit of this. We can leverage xtensor and the checks that come from pymc, xtensor, pymc_extras, etc. Here are some items:

All 6 helper functions are redundant with existing APIs.

_supports_xdist (inspect.signature check) — Prior, Censored, and VariableFactory all accept xdist: bool = False. pymc_extras handles the introspection internally. Just pass xdist=True.

_exp — ptx.math.exp works on floats, xtensor, and xtensor scalars. No dispatch needed.

_create_dim_variable — prior.create_variable(name, xdist=True) is one line.

_align_to_dims — pmd.Deterministic(name, value, dims=...) already broadcasts and transposes.

_check_dims_known — pmd.Deterministic / pmd.Data raise ValueError for missing dims.

_create_likelihood_variable — Prior.create_likelihood_variable already has all guards (UnsupportedDistributionError, MuAlreadyExistsError). For observed: prior.create_likelihood_variable(xdist=True). For unobserved: attach mu manually + create_variable (recommended workaround for pymc-extras#731, future-proof for upcoming deprecation).

The MMM codebase uses Prior.create_variable(xdist=True), ptx.math.*, and pmd.Deterministic with dims= throughout (transformers.py, link.py, fourier.py, additive_effect.py). The Bass model should match.

F() and f() need ptx.math.exp and ptx.math.square — pt.exp / pt.square raise TypeError on xtensor variables.

test_derivative — pytensor.grad cannot work on xtensor graphs. @pytest.mark.xfail until pytensor#2337 merges.

~6 tests in TestBassModelXdistFallbacks exercise removed code paths and should be dropped.

combined_dims declaration-order fix, peak dims change, core_op unwrap in mlflow, and all TestBassModelDims tests are correct and should stay.

[anevolbap]

Thanks for the detailed pass, it made me run every item against the released stack (pymc 6.2.0, pymc-extras 0.14.0). Good news first: the observed path does collapse to create_likelihood_variable(xdist=True), _check_dims_known partly goes: 6.2 raises ValueError at dims= registration (m/p/q), though the observed path and the deterministic broadcast still need it, the bare failure there is a KeyError from inside pymc, and the unobserved branch is already the #731 mu-attach workaround. Three items break when run, though:

• Prior("Wald", mu=10, lam=1).create_variable("w", xdist=True) raises UnsupportedDistributionError (nested case too), and a VariableFactory without the kwarg raises TypeError. 0.14 has no internal fallback; that is what _supports_xdist covers.
• ptx.math.exp raises TypeError on numpy arrays and dim-less tensor vectors, and pytensor.grad fails on the result. _exp keeps the public F/f working on the inputs they accept today, which also keeps test_derivative green with no xfail.
• pmd.Deterministic(dims=...) transposes but does not broadcast: a value on ("T",) with dims=("T", "product") raises ValueError. The broadcast is what _align_to_dims adds.

So the full cleanup means narrowing the contract to pymc.dims-supported distributions and xdist-aware factories, same as MMM. Is that the intent? If yes, the fallback helpers and the ~6 tests go and Wald-style priors stop working; happy to push that version.

[williambdean]

We don't need to support everything. We want to leverage xtensor mainly like with the MMM implementation.
No need to "broadcast" a Deterministic. If it doesn't have that dimension, we can leave it to the user to do what they would like with it outside of the model.

The model block should be lean.

[anevolbap]

Pushed the lean version (42039a8): _supports_xdist, _create_dim_variable, _check_dims_known and _align_to_dims are gone, m/p/q and the likelihood go straight through create_variable(xdist=True), and the deterministics are no longer broadcast. 108 bass tests green, model block down ~90 lines.

Two things stayed, with the reason:

• _exp in F/f. ptx.math.exp raises TypeError: Cannot convert [...] to XTensorType on a plain ndarray, so f(0.03, 0.38, np.arange(5)) and every fixture built on it stops working, and pytensor.grad then fails on the xtensor graph. Keeping the dispatch is also what avoids the test_derivative xfail.
• The deterministics still get a transpose into (T, ...). Without it m * f(p, q, t) comes out as ('product', 'T') and the stored curves flip layout. No broadcasting, just the order.

One tradeoff to flag: without _check_dims_known, an unknown dim on observed now surfaces as KeyError: 'geo' from inside pymc_extras instead of a message naming the coord. The likelihood-dims case is unaffected, pymc's own error there is already clear. Happy to inline a two-line guard if that message is worth keeping.

[williambdean]

Why the _exp 🥲

Just update the fixtures 😅

[anevolbap]

As Claude usually says, "You're absolutely right!" — _exp is gone (cbaed2f).

On test_derivative: skipped the xfail. Lowering the graph first (rewrite_graph(..., include=("lower_xtensor",))) lets pytensor.grad walk it, so the dF/dt against f check still runs instead of going dormant until pytensor#2337 lands.

[daimon-pymclabs]

Review notes. Two of these I'd want closed before merge, two are cheap and belong in this PR rather than a follow-up. Everything else I found is listed at the bottom as non-blocking.

---

1. copy.deepcopy(priors["likelihood"]) detaches tensor-valued parameters (blocker)

model.py:409. The comment above it is right about the problem (mutating the caller's prior leaks this model's dims into a reused config) but deepcopy is a heavier instrument than the problem needs, and it has a side effect.

Prior.__deepcopy__ does copy.deepcopy(self.parameters), so any pytensor variable sitting in a parameter is cloned. Confirmed on pytensor 3.3.0:

>>> x = pt.scalar("x"); copy.deepcopy(x) is x
False
>>> copy.deepcopy(x * 2).owner.inputs[0] is x
False

So for a likelihood like Prior("Normal", sigma=some_pm_data), the y that gets built is wired to a clone of some_pm_data that is not in the model. pm.set_data on the original then silently stops affecting y — no error, just a posterior that ignores the update. Prior.deepcopy() is not an escape hatch, it is copy.deepcopy(self).

The mutation being guarded against is exactly one attribute, so scope it to that attribute. Censored.dims is a property whose setter forwards to self.distribution.dims, so the same code covers both types:

from contextlib import contextmanager

@contextmanager
def _borrow_dims(prior: Prior | Censored, dims: tuple[str, ...]):
    """Set ``dims`` for the duration of the block, leaving the caller's prior as found.

    A copy is not usable here: ``Prior.__deepcopy__`` deep-copies ``parameters``,
    which clones any pytensor variable it finds and detaches it from the model
    graph, so a later ``set_data`` would no longer reach the likelihood.
    """
    original = prior.dims
    prior.dims = dims
    try:
        yield prior
    finally:
        prior.dims = original

and at the call site:

with _borrow_dims(priors["likelihood"], combined_dims) as likelihood:
    observed_xt = ...  # unchanged
    _create_likelihood_variable(likelihood, "y", mu=adopters, observed=observed_xt)

test_likelihood_prior_is_untouched passes either way, which is why it did not catch this. The regression test that would:

def test_shared_likelihood_parameter_stays_connected():
    """A tensor in a likelihood parameter must remain the model's own node."""
    with pm.Model(coords={"T": np.arange(5)}) as model:
        sigma = pm.Data("sigma", 1.0)
        create_bass_model(
            t=np.arange(5),
            observed=np.arange(5),
            priors={..., "likelihood": Prior("Normal", sigma=sigma, dims="T")},
            coords={"T": np.arange(5)},
            model=model,
        )
    assert sigma in pytensor.graph.basic.ancestors([model["y"]])

Same hazard applies to inner.deepcopy() at model.py:284. There the copy is genuinely needed (you are adding mu to parameters), so it cannot just be dropped — but it means the observed=None path detaches tensor parameters too. Minimum viable: copy the container and rebind the original parameter objects,

unobserved = inner.deepcopy()
unobserved.parameters = {**inner.parameters, "mu": mu}

which keeps every original parameter object identical while still leaving inner untouched.

2. The observed dim lookup can pick up an unrelated variable's dims (blocker)

model.py:417-420:

observed_dims = getattr(observed, "dims", None) or (
    model.named_vars_to_dims.get(getattr(observed, "name", None), combined_dims)
)

By the time this runs, named_vars_to_dims holds adopters, innovators, imitators, peak, m, p, q. A plain pt.TensorVariable that happens to be named peak — not registered with the model, just named — gets labelled with peak's dims and transposed without complaint. Name equality is being used as a proxy for identity.

Also, pm.Data can be registered with a partially-None dims tuple, and pmd.as_xtensor(obs, dims=(None, "product")) fails from deep inside with an error that does not point back here.

Both are covered by one helper:

def _observed_dims(observed, model: Model, combined_dims: tuple[str, ...]) -> tuple[str, ...]:
    """Axis labels for ``observed``: its own, else the model's, else positional."""
    own = getattr(observed, "dims", None)
    if own:
        return tuple(own)

    name = getattr(observed, "name", None)
    # Name equality is not identity: an unregistered variable may share a name
    # with a model variable, and would otherwise borrow its dims.
    if name is not None and model.named_vars.get(name) is observed:
        registered = model.named_vars_to_dims.get(name)
        if registered and all(dim is not None for dim in registered):
            return tuple(registered)

    return combined_dims

Test:

def test_unregistered_variable_sharing_a_name_is_labelled_positionally():
    """Name equality must not be mistaken for identity."""
    # a bare variable named "peak", never registered on the model
    ...
    assert model["y"].dims == ("T", "product")  # not peak's dims

3. F / f lost plain-array support with no deprecation path (should fix here)

These are public and documented as array-like. F(p, q, np.arange(10)) worked on main; now every call site needs pmd.as_xtensor, and every test in this diff was updated to do exactly that. The module docstring records the change, but nothing in the code does, and there is no test pinning what a user hitting this actually sees.

The signatures also disagree with the tests: t: XTensorVariable, but test_derivative still passes pt.scalar.

Two options, either is fine, but pick one explicitly:

a. Accept array-like and wrap — no break, one line each:

def _as_time(t) -> XTensorVariable:
    """Accept an xtensor, or label a 1-d array-like positionally as ``("T",)``."""
    if isinstance(t, XTensorVariable):
        return t
    return pmd.as_xtensor(pt.as_tensor(t), dims=("T",))

b. Keep the break and make it legible — raise at the boundary and test the message:

if not isinstance(t, XTensorVariable):
    raise TypeError(
        f"`t` must be an XTensorVariable, got {type(t).__name__}. "
        "Wrap plain arrays with `pymc.dims.as_xtensor(t, dims=('T',))`."
    )

plus pytest.raises(TypeError, match="as_xtensor"), and fix the annotations so the 0-d/scalar cases the tests rely on are actually covered by the declared type.

4. The deterministics dim change needs to be in the release notes (should fix here)

Dropping the broadcast so pooled p/q/m produce pooled curves is the right call and is well tested, but it changes posterior shapes and turns plot_* output for existing multi-series users from N subplots into one. The PR body says so; nothing in the repo does. With no CHANGELOG file, that means the release notes explicitly, along the lines of:

Breaking: adopters, innovators, imitators and peak now carry only the dims their parameters carry. Previously they were broadcast to all combined dims, so a pooled p/q/m produced N identical curves; they now produce one. Posterior shapes and plot_* subplot counts change accordingly for models where any of p, q, m is pooled. Custom VariableFactory implementations for m/p/q must now accept an xdist keyword.

Worth an eyeball on the bass example notebook to confirm its narrative still matches the new shapes.

---

Non-blocking

• _create_likelihood_variable reimplements pymc-extras internals (signature(inner.pymc_distribution.dist).parameters). The docstring cites pymc-extras#731; add a # TODO(pymc-extras#731): drop once observed=None is supported upstream at the code so it gets deleted rather than fossilised.
• observed: pt.TensorLike | None (model.py:305) should include xr.DataArray in the annotation and docstring — several new tests depend on that path.
• The deterministic() closure builds order by filtering combined_dims, so a dim on value that is not in combined_dims is dropped from the label rather than raising. Unreachable today; assert set(value.dims) <= set(combined_dims) keeps it that way.
• BassPriors types m/p/q as VariableFactory, but factories must now accept xdist. One that does not gets a bare TypeError from deep inside; a checked error naming the parameter would be kinder.
• Plotting error text: when the user passes an explicit unknown dim, "give p, q or m that dim" is the wrong advice for what is usually a typo'd dim name.
• Test gaps: a Censored likelihood with extra dims and observed=None never exercises dim propagation through the reconstructed Censored(unobserved, ...) — the existing test only covers combined_dims == ("T",). And there is no coverage of the custom-VariableFactory path this PR flags as breaking.
• lower() depends on the internal rewrite name "lower_xtensor". Annotated with pytensor#2337, which is the right thing to do; just noting it is a fragile hook.

Nice work overall — the migration is clean, and test_observed_is_labelled_with_its_own_dims and test_combined_dims_follow_declaration_order are exactly the right regression guards for the two subtle bugs this touches.