Laplace approximation not handling non-scalar parameters

[dar326]

I have been trying to utilize the new laplace function for estimating the posterior distribution using a quadratic approximation. It seems that there is an issue when trying to use non-scalar parameters. This is a slight modification of the example that is provided in the documentation for the function:

y = np.array([2642, 3503, 4358]*10)
cat = np.array([0]*15 + [1]*15)
with pm.Model() as m:
    logsigma = pm.Uniform("logsigma", 1, 100)
    mu = pm.Uniform("mu", -10000, 10000, shape=2)
    yobs = pm.Normal("y", mu=mu[cat], sigma=pm.math.exp(logsigma), observed=y)
    idata = laplace([mu, logsigma], model=m)

The following error is generated:

TypeError: ('Wrong number of dimensions: expected 1, got 0 with shape ().', 'Container name "mu"')

Interestingly, this error occurs even when the shape is not greater than 1 and no indexing is performed:

y = np.array([2642, 3503, 4358]*10)
with pm.Model() as m:
    logsigma = pm.Uniform("logsigma", 1, 100)
    mu = pm.Uniform("mu", -10000, 10000, shape=1)
    yobs = pm.Normal("y", mu=mu, sigma=pm.math.exp(logsigma), observed=y)
    idata = laplace([mu, logsigma], model=m)

This seems to indicate that there is an issue with setting the shape parameter to any value.

[ricardoV94]

I can confirm this is a bug. The code assumes all parameters to be scalars in multiple places. First this line would fail with 2d parameters:

pymc-experimental/pymc_experimental/inference/laplace.py

Line 120 in f61e161

mean = np.concatenate([np.atleast_1d(map[v.name]) for v in vars])

Which should be instead mean = np.concatenate([np.atleast_1d(map[v.name].ravel()) for v in vars])

Then in this loop we are assuming there is a single entry per variable (with chain, draws batch):

pymc-experimental/pymc_experimental/inference/laplace.py

Line 129 in f61e161

data_vars[str(var)] = xr.DataArray(samples[:, :, i], dims=("chain", "draw"))

This later leads to the error reported above.

Instead we should slice the portion corresponding to each variable, and reshape it to the original shape. There should be some utilities in the PyMC codebase to help with that. As a bonus we may want to recover the core dims, although I think the code would work without them.

[ricardoV94]

CC @carsten-j

[carsten-j]

@ricardoV94 I will try to take at look at this sometime next week.

[carsten-j]

@ricardoV94 I am not sure that I understand how to fix this bug.

[jessegrabowski]

I think this can be fixed by the new code in #385