New (experimental) way of writing PyMC models that leverages dims

Announcements
1

In the 5.25 release of PyMC we introduced a new dims module, that let’s you write PyMC models differently.

If the traditional API looks like writing numpy code, the new one looks like writing xarray code.

Here is a redacted example from the new documentation notebook

import pymc as pm
import pymc.dims as pmd  # <- NEW

# ... create some fake-data and coords

with pm.Model(coords=coords) as dims_splines_model:
    x = pmd.Data("x", x_np, dims="obs")
    y_obs = pmd.Data("y_obs", y_obs_np, dims="obs")
    knots = pmd.Data("knots", knots_np, dims=("knot",))
    group_idx = pmd.Data(group_idx_np, dims=("obs",))

    sigma = pmd.HalfCauchy("sigma", beta=1)
    sigma_beta0 = pmd.HalfNormal("sigma_beta0", sigma=10)
    beta0 = pmd.HalfNormal("beta_0", sigma=sigma_beta0, dims=("group",))
    z = pmd.Normal("z", dims=("group", "knot"))

    delta_factors = pmd.math.softmax(z, dim="knot")
    slope_factors = 1 - delta_factors.isel(knot=slice(None, -1)).cumsum("knot")
    spline_slopes = pmd.concat([beta0, beta0 * slope_factors], dim="knot")
    beta = pmd.Deterministic("beta", pmd.concat([beta0, spline_slopes.diff("knot")], dim="knot"))

    X = pmd.math.maximum(0, x - knots)
    mu = (X * beta.isel(group=group_idx)).sum("knot")
    y = pmd.Normal("y", mu=mu, sigma=sigma, observed=y_obs)

It has a lot going on, but the gist is that axis arguments are replaced by dims arguments, and one no longer needs to keep track of dimension order, or align them manually with transposes, np.newaxis/None and empty slices :, everywhere.

If two variables going into a function have the same dimensions they are automatically aligned, while distinct dimensions broadcast.

Everything is explained in more detail in the documentation notebook.

This is a work-in-progress, several distributions are still missing and we may change the API drastically, so don’t use it for serious work (just yet).

While we iterate, we would love feedback from the braver users! Let us know if you find any bug in our github, or there’s some functionality you absolutely need to get started: GitHub · Where software is built

10 Likes