predict fails with Variational Inference

[omrihar]

I'm trying to use predict on out of sample data after sampling using variational inference. When I try to predict I get an "unexpected dimensions" error.

Minimal code to reproduce:

import pandas as pd
import bambi as bmb

test_model = bmb.Model('y ~ x', pd.DataFrame({'x': [1,2,3], 'y': [2, 3, 4]}))
test_approx = test_model.fit(inference_method='vi')
test_tr = test_approx.sample(draws=1000)

test_model.predict(idata=test_tr, data=pd.DataFrame({'x': [4, 5], 'y': [1,1]}))

[tomicapretto]

There's definitely something going on. The predictor x has a dimension when it shouldn't.

I tried reproducing the issue with PyMC, and I couldn't. So I realized it is something about how Bambi builds the model.

Then I looked into the random variables in the underlying PyMC model and found the problem:

test_model.backend.model.free_RVs
# [Intercept ~ N(3, 5.4), x ~ N(0, 2.5), y_sigma ~ HalfStudentT(4, 0.816)]
test_model.backend.model.free_RVs[1].shape
# TensorConstant{(1,) of 1}

which means that the distribution assigned for the slope is 1 dimensional when it should be scalar. That's why it adds an extra dimension.

If we have a look at the PyMC model, this does not happen

import pymc as pm

x = [1, 2, 3]
y = [2, 3, 4]

with pm.Model() as pm_model:
    a = pm.Normal("a")
    b = pm.Normal("b")
    pm.Normal("y", mu=a + b * x, sigma=pm.HalfNormal("sigma"), observed=y)

pm_model.free_RVs
# [a ~ N(0, 1), b ~ N(0, 1), sigma ~ N**+(0, 1)]
pm_model.free_RVs[1].shape
# TensorConstant{[]}

[omrihar]

Is there a way to circumvent this until the PR gets merged? For example, changing the dimensionality of the scalars after fitting, so that the predict method would work?

[tomicapretto]

I don't know. But if you don't want to wait until the PR is merged, you can monkey-patch Bambi locally changing only these lines

shape = None if data.shape[1] == 1 else data.shape[1]
coef = distribution(label, shape=shape, **args)
coef = at.atleast_1d(coef)

in line 51 of bambi/backend/terms.py.

See the screenshot

All the other changes are related to predictions

[tomicapretto]

Since #575 is merged, we can close this one. Thanks for reporting the issue @omrihar!