Implement MultiDimensionalMMM with pymc.dims

[ricardoV94]

TODO:

• Dev migration guide
• Adapt new incrementality code
• Censored Prior

Closes #1981
Closes #1514
Closes #1630 (just a regression test, it worked fine before)

---

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

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[twiecki]

That's awesome.

[codecov]

Codecov Report

❌ Patch coverage is 95.63025% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.52%. Comparing base (26ef53f) to head (3d37c19).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc_marketing/mmm/utility.py	81.25%	9 Missing ⚠️
pymc_marketing/mmm/components/base.py	89.47%	4 Missing ⚠️
pymc_marketing/special_priors.py	88.88%	4 Missing ⚠️
pymc_marketing/mmm/hsgp.py	94.54%	3 Missing ⚠️
pymc_marketing/mmm/incrementality.py	0.00%	3 Missing ⚠️
pymc_marketing/mmm/sensitivity_analysis.py	94.11%	2 Missing ⚠️
pymc_marketing/mmm/linear_trend.py	95.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2204      +/-   ##
==========================================
- Coverage   93.20%   88.52%   -4.69%     
==========================================
  Files          79       80       +1     
  Lines       12477    12485       +8     
==========================================
- Hits        11629    11052     -577     
- Misses        848     1433     +585     

☔ View full report in Codecov by Sentry.
📢 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.

[cetagostini]

Hey! I found a few typos while reviewing the pymc.dims implementation. See inline comments below 👇

[cetagostini]

A few more inconsistencies I noticed - some return type annotations say TensorVariable but the code actually returns XTensorVariable 🤔

[cetagostini]

Function Equivalence Testing Results 🧪

I tested the transformation functions to verify they produce equivalent outputs with the new pymc.dims API. Here's what I found:

✅ All Functions Pass Numerical Equivalence

Function	Match
logistic_saturation	✅
tanh_saturation	✅
hill_function	✅
michaelis_menten	✅
root_saturation	✅
geometric_adstock weights	✅

⚠️ Important Finding: Dimension Ordering

When broadcasting with xtensor, the dimension order may differ from numpy!

# Example:
x_2d = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])  # (3 dates, 2 channels)
lam_channel = np.array([0.3, 0.7])  # per-channel lambda

# Old API (numpy-style):
# Output shape: (3, 2) → (date, channel) order preserved

# New API (xtensor):
# Output shape: (2, 3) → (channel, date) order CHANGED!

The values are correct (match after transpose), but any code relying on axis indices (axis=0, axis=1) instead of dimension names may break.

This is expected behavior for xarray-like semantics, but worth noting for anyone working with the codebase! 👍

[ricardoV94]

Yes order doesn't matter while in xtensor land. You can always transpose explicitly to another order, which you should do if you need to go to tensor land (via .values).

Note even if not so relevant, dimension order matches the same as xarray

[cetagostini]

This is very nice!

[ricardoV94]

Legacy MMM/Optimizer now work with pymc.dims and the refactored components.

This PR now depends on pymc-devs/pymc-extras#657 to help with the transition into dim world, so the CI will be ready until we re-sync

[ricardoV94]

Failing test depends on pymc-devs/pymc#8124
Notebook failing due to prior.Censored not being implented with pymc.dims yet