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

[ricardoV94]

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]

wtf 😆

[ricardoV94]

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]

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

[ricardoV94]

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]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[twiecki]

We should make this the 6.0 release.

[ricardoV94]

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