terms#

Composable model terms for building pymc.dims subgraphs from xarray data.

pymc_marketing.terms provides a small set of built-in pieces (Parameter / Intercept, Dot, Transform, Named, Ref) that compose with +, *, and - into predictors. This module is not a full modeling toolkit — it does not ship complete hierarchical, domain-specific, or observation-model wrappers. It is a thin, expressive core that lets you build what you need by subclassing ModelTerm and reusing the same helpers and operators that built-ins use. If a term is missing from this package, that is the intended path: write a custom term, compose it with + / * / -, and the framework handles coords, data registration, tensor construction, and prediction updates.

The framework is pymc.dims-native. create_variable() returns an XTensorVariable (dimensional tensor). Built-in terms use pmd.Data for shared data and pmd.Normal with xdist=True for variables. Custom terms should use pymc.dims and pytensor.xtensor throughout to ensure compatibility during composition.

Extending#

This module is designed so that the parts you do not find here are the parts you are meant to write. ModelTerm is the extension base; subclass it, implement five lifecycle methods, and the result composes with the built-ins via + / * / - — no framework changes required.

Lifecycle (per term)#

  • get_coords(ds) – return coordinates from data variables (outside model)

  • add_coords(ds) – add model-only coordinates, if not in dataset (inside model context)

  • register_data(ds) – register pmd.Data shared variables (inside model)

  • create_variable() – build a dimensional tensor inside a pm.Model context

  • set_data(ds, model) – update shared variables for prediction

register_data calls add_coords before register_data on each ModelTerm, so dynamic coordinates are handled automatically.

Rules#

  • get_coords should be pure (no model context).

  • Guard pmd.Data registration with if name not in model to share data between terms.

  • create_variable must return a dims-compatible (xtensor) tensor.

  • Each RV name must be unique across the full expression tree.

  • build_param accepts VariableFactory, ModelTerm, xr.DataArray, or numeric literal. Prior.create_variable is called with xdist=True.

Serialization#

Built-in terms serialize through pymc_marketing.serialization (to_dict / from_dict are registered), so recipes stored in model_config survive fit() (attrs are JSON-serialized at sampling time) and save() / load().

Custom terms do not need subclassing to participate (the protocol stays open); serialization has two routes:

  • Dataclass terms are the low-ceremony route: because all term containers are dataclasses, this is the recommended form for custom terms (see pymc_marketing.serialization).

  • Non-dataclass terms implement to_dict / from_dict and register with @serialization.register (pydantic classes can inherit SerializableBaseModel to get both automatically). Serializing a composition that contains an unknown term raises SerializationError.

Children of composed terms (Sum, Product, Transform) may be terms, VariableFactory (Prior, …), xr.DataArray, DeferredFactory, or numeric literals. Transform functions serialize by name and are resolved with pymc_extras.prior._get_transform: pytensor.xtensor.math / pytensor.xtensor.linalg functions, or transforms registered with pymc_extras.prior.register_tensor_transform.

Parameter and Intercept#

These are the same kind of term: a named free parameter whose shape comes from the prior (scalar when no dims; dimensional when the prior carries dims). Intercept is a convenience alias with name defaulting to "intercept"; use Parameter when the role is not an intercept (e.g. a cohort scale or product effect). When the prior carries dims, coordinates are automatically extracted from the dataset during collect_coords.

Example: custom gather term#

Copy this into your project as a starting point for cohort / index / group gathers:

@dataclass
class Indexed(ModelTerm):
    base: Any
    index_var: str

    def get_coords(self, ds):
        return get_coords(self.base, ds)

    def register_data(self, ds):
        register_data(self.base, ds=ds)
        model = pm.modelcontext(None)
        if self.index_var not in model:
            pmd.Data(
                self.index_var, ds[self.index_var], dims=ds[self.index_var].dims
            )

    def create_variable(self):
        base_val = build_param(self.base)
        model = pm.modelcontext(None)
        idx = model[self.index_var]
        return base_val[idx]

    def set_data(self, ds, model=None):
        set_data(self.base, ds=ds, model=model)
        if self.index_var in ds:
            pm.set_data({self.index_var: ds[self.index_var].values}, model=model)

Usage (custom + built-ins composed in one expression):

import pytensor.xtensor as ptx
from pymc_extras.prior import Prior
from pymc_marketing.terms import (
    Dot,
    Parameter,
    Transform,
    build_param,
    collect_coords,
    get_coords,
    register_data,
    set_data,
)

mu = Indexed(
    Parameter("alpha_scale", prior=Prior("HalfNormal", dims="cohort")),
    "cohort_idx",
) * Transform(
    -Dot(var_name="covars", name="coef_alpha", prior=Prior("Normal", dims="covar")),
    ptx.math.exp,
)
coords = collect_coords(mu, ds=ds)
with pm.Model(coords=coords) as model:
    register_data(mu, ds=ds)
    alpha = build_param(mu)

Examples#

Built-ins cover the common pieces. The explicit four-step lifecycle (collect_coords / register_data / build_param / set_data) works the same for built-ins and custom terms — no hidden helpers.

import pymc.dims as pmd
import pytensor.xtensor as ptx
from pymc_extras.prior import Prior
from pymc_marketing.terms import (
    Dot,
    Intercept,
    Transform,
    build_param,
    collect_coords,
    register_data,
)

mu = Intercept(name="mu") + Dot(var_name="x", prior=Prior("Normal", dims="feature"))
sigma = Transform(Intercept(name="sigma"), func=ptx.math.exp)

coords = collect_coords(mu, sigma, ds=ds)

with pm.Model(coords=coords) as model:
    register_data(mu, ds=ds)
    register_data(sigma, ds=ds)
    pmd.Normal(
        "y_obs", mu=build_param(mu), sigma=build_param(sigma), observed=ds["y"]
    )

from pymc_marketing.terms import (
    Parameter,
    Dot,
    collect_coords,
    register_data,
    build_param,
)

mu = Parameter("baseline", prior=Prior("Normal", dims="product")) + Dot(
    var_name="x", prior=Prior("Normal", dims="feature")
)
coords = collect_coords(mu, ds=ds)
with pm.Model(coords=coords) as model:
    register_data(mu, ds=ds)
    build_param(mu)

from pymc_marketing.terms import (
    Dot,
    Intercept,
    collect_coords,
    register_data,
    build_param,
)

alpha_branch = Dot(
    var_name="x", name="alpha_coef", prior=Prior("Normal", dims="feature")
)
beta_branch = Dot(
    var_name="x", name="beta_coef", prior=Prior("Normal", dims="feature")
)
mu = alpha_branch + Intercept(name="offset")
# beta_branch can be used in a separate expression tree

Gotchas#

  • create_variable() should return a dimensional tensor (via pmd.Normal with xdist=True, or pytensor.xtensor). Mixing non-xtensor output (e.g., regular pm.Normal) with xtensor output in a Sum will fail when build_param composes them.

  • Two Dot terms referencing the same var_name will share the pmd.Data variable automatically (duplicate registration is skipped).

  • Two Dot terms with the same var_name will also create identically named beta variables ({var_name}_beta), causing a PyMC “already exists” error. Pass a distinct name to each Dot to give the coefficients different variable names while still sharing the same pmd.Data (e.g. separate alpha and beta coefficient branches on the same covariate matrix).

  • build_param is not idempotent — call it once to create PyMC variables, then sample. Calling it a second time tries to create variables with the same name, causing a PyMC error. Named is a named wrapper with the same constraint: build once, reference afterwards with Ref (which is freely reusable inside multiple Named expressions).

  • Changing the number of observations in set_data requires rebuilding the model. This is a pmd.Data constraint (dimensional shared variables have fixed dimension sizes), not a framework limitation.

  • Model-internal coordinates that are not in the dataset can be added by overriding add_coords:

    def add_coords(self, ds):
    model = pm.modelcontext(None)
    unique = list(dict.fromkeys(ds[self.data_source].values))
    model.add_coords({self.data_source: unique})

  • Serialization contract: see the “Serialization” section above.

  • Named builds a pmd.Deterministic, so its expression leaves must be dims-native: Parameter / Prior with xdist=True or pytensor.xtensor expressions. Plain-tensor leaves (xdist=False) will fail when the deterministic is built.

Functions

build_param(param[, name])

Build a tensor parameter from a term, composition, variable factory, or constant.

collect_coords(*params, ds)

Collect coordinates from multiple parameter expressions.

collect_terms(params)

Collect all ModelTerm instances from a list of parameter specifications.

data_vars(param)

List the dataset variables a term tree reads, at any depth.

frozen_deterministics(param)

Collect the frozen-deterministic names a term tree declares, at any depth.

get_coords(param, ds)

Extract coordinates from a term or term composition.

register_data(param, *, ds)

Register shared data variables as pmd.Data in the active model.

set_data(param, *, ds, model)

Update shared data variables for out-of-sample prediction.

Classes

Dot(*, var_name, prior[, name])

Linear predictor term: data @ beta.

Intercept([name, prior])

GLM-style intercept term; same as Parameter with a default name.

ModelTerm()

Base class for composable model terms.

Named(name, expr[, dims])

Compose an expression into a named deterministic variable.

Parameter(name[, prior])

Named free parameter (scalar or dimensional via prior.dims).

Product(left, right)

Container for multiplicative composition via *.

Ref(name)

Reference an already-built named variable in the active model.

Sum(terms)

Container for additive composition via +.

Transform(inner, func)

Apply a pytensor function to a term's output.