Laplace approximation TypeError

Hello. I would like to use the Laplace Approximation in pymc-extras. Both MCMC and ADVI work fine, but LA throws a TypeError as shown below. The source code is too complicated for me to figure out the cause. Below are my model and the error message. What is causing this error?

Model

from patsy import dmatrix

# blossom_data = pd.read_csv(pm.get_data("cherry_blossoms.csv"), sep=";")
blossom_data = pd.read_csv('./data/cherry_blossoms.csv')
blossom_data = blossom_data.dropna(subset=["doy"]).reset_index(drop=True)

num_knots = 15
knot_list = np.percentile(blossom_data.year, np.linspace(0, 100, num_knots + 2))[1:-1]
B_cov = dmatrix(
    "bs(year, knots=knots, degree=3, include_intercept=True) - 1",
    {"year": blossom_data.year.values, "knots": knot_list},
)
B_cov = np.asarray(B_cov)
Y = blossom_data.doy.values

COORDS = {"splines": np.arange(B_cov.shape[1])}
with pm.Model(coords=COORDS) as pymc_model:
    obs = pm.Data('obs', Y)
    B = pm.Data('B', B_cov)

    a = pm.Normal("a", 100, 5, shape=(1,))
    w = pm.Normal("w", mu=0, sigma=3, dims="splines")

    mu = pm.Deterministic("mu", a + pt.dot(B, w))
    sigma = pm.Exponential("sigma", 1, shape=(1,))

    _ = pm.Normal("y", mu=mu, sigma=sigma, observed=obs)

Run Code

with pymc_model:
    idata = pmx.fit(
        method="laplace",
        chains=4,
        draws=2500,
    )

Error Message

File ~/.conda/envs/pymc_env3/lib/python3.13/site-packages/pymc_extras/inference/laplace_approx/laplace.py:223, in model_to_laplace_approx(model, unpacked_variable_names, chains, draws)
    220             initval = initial_point.get(name, None)
    221             dim_shapes = initval.shape if initval is not None else batched_rv.type.shape[2:]
    222             laplace_model.add_coords(
--> 223                 {name: np.arange(shape) for name, shape in zip(dims[2:], dim_shapes)}
    224             )
    226         pm.Deterministic(name, batched_rv, dims=dims)
    228 return laplace_model

TypeError: unsupported operand type(s) for -: 'NoneType' and 'int'

This is a bug. I opened an issue about it here. You can side-step it by setting coords on mu:

COORDS = {"splines": np.arange(B_cov.shape[1]), 
          'obs_idx': np.arange(Y.shape[0])}

with pm.Model(coords=COORDS) as pymc_model:
    obs = pm.Data('obs', Y, dims=['obs_idx'])
    B = pm.Data('B', B_cov)

    a = pm.Normal("a", 100, 5, shape=(1,))
    w = pm.Normal("w", mu=0, sigma=3, dims="splines")

    mu = pm.Deterministic("mu", a + pt.dot(B, w), dims=['obs_idx'])
    sigma = pm.Exponential("sigma", 1, shape=(1,))

    _ = pm.Normal("y", mu=mu, sigma=sigma, observed=obs)
1 Like

Thank you so much!