Add experimental `dims` module with objects that follow dim-based semantics (like xarray) #7820

ricardoV94 · 2025-06-17T18:12:40Z

This builds on top of PyTensor xtensor module, to introduce distributions and model objects that follow xarray-like semantics. Example model:

import numpy as np
import pymc as pm
import pymc.dims as pmd

# Very realistic looking data!
observed_response_np = np.ones((5, 20), dtype=int)
coords = coords = {
    "participant": range(5),
    "trial": range(20),
    "item": range(3),
}
with pm.Model(coords=coords) as dmodel:
    observed_response = pmd.Data(
        "observed_response", observed_response_np, dims=("participant", "trial")
    )

    # Participant constant preferences for each item
    participant_preference = pmd.ZeroSumNormal(
        "participant_preference", 
        core_dims="item", 
        dims=("participant", "item"),
    )

    # Shared time effects across all participants
    time_effects = pmd.Normal("time_effects", dims=("item", "trial"))

    trial_preference = pmd.Deterministic(
        "trial_pereference",
        participant_preference + time_effects,
        dims=(...,),  # No need to specify, PyMC knows them
    )

    response = pmd.Categorical(
        "response",
        p=pmd.math.softmax(trial_preference, dim="item"),
        core_dims="item",
        observed=observed_response,
        dims=(...,), # No need to specify, PyMC knows them
    )

Equivalently, with the traditional API:

with pm.Model(coords=coords) as model:
    observed_response = pm.Data(
        "observed_response", observed_response_np, dims=("participant", "trial")
    )

    # Participant constant preferences for each item
    participant_preference = pm.ZeroSumNormal(
        "participant_preference", 
        n_zerosum_axes=1,
       dims=("participant", "item"),
    )

    # Shared time effects across all participants
    time_effects = pm.Normal("time_effects", dims=("trial", "item"))

    trial_preference = pm.Deterministic(
        "trial_preference",
        participant_preference[:, None, :] + time_effects[None, :, :],
        dims=("participant", "trial", "item"),
    )

    response = pm.Categorical(
        "response",
        p=pm.math.softmax(trial_preference, axis=-1),
        observed=observed_response,
        dims=("participant", "trial"),
    )

More details in the new core notebook

twiecki · 2025-06-18T00:15:03Z

pymc/dims/__init__.py

+    day_of_conception = datetime.date(2025, 6, 17)
+    day_of_last_bug = datetime.date(2025, 6, 17)
+    today = datetime.date.today()
+    days_with_bugs = (day_of_last_bug - day_of_conception).days


This has two purposes: distract reviewers so they don't focus on the critical changes, and prove that OSS libraries can't be fun.

twiecki · 2025-06-21T16:34:32Z

Can this index using labels? x["a"]

ricardoV94 · 2025-06-21T16:54:40Z

Can this index using labels? x["a"]

I don't know what x["a"] means :).

Is "a" a coordinate? x.loc["a"] would be the xarray syntax? You can't do that.

Like in xarray, you can do x.isel(dim=idxs) or x[{dim: idxs}].

You cannot do x.sel(dim=coords) or x.loc[coords]

The new PyTensor objects have dims but not coords. It's not trivial to encode coord based operations in our backends.

review-notebook-app · 2025-06-24T18:07:52Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

twiecki · 2025-06-30T14:47:24Z

We should make this the 6.0 release.

ricardoV94 · 2025-06-30T15:00:29Z

We should make this the 6.0 release.

I agree, but would perhaps wait until we beta-tested it to the point it no longer feels too experimental

OriolAbril

adding one extra comment about myst syntax but outside reviewnb because it always messes up rendering of backticks.

I will try to play around and build/port some models one of these days, and look at the code itself while I do that

OriolAbril · 2025-07-02T11:37:10Z