Can't get predictions from Bambi. Encountering `ValueError: different number of dimensions on data and dims: 0 vs 2`

(apologies this is a bit of a cross post — I posted as a bambi github issue first before it occurred to me that pymc may have a more populated forum that might also be a better place since I’m not sure if this is user error or a bug)

Hi All!

New to Bambi, and really struggling with predicting. Can’t tell if I’m doing something wrong or found a bug.

When I try to predict new values of a simple hierarchical model, I’m getting a ValueError: different number of dimensions on data and dims: 0 vs 2 error I can’t figure out. Any chance of help? Thanks!

Stack trace down below.

Versions:

bambi                      0.16.0           pyhcf101f3_1     conda-forge
pymc                       5.28.0           hc029463_0       conda-forge
pymc-base                  5.28.0           pyhcf101f3_0     conda-forge
arviz                      0.23.4           pyhcf101f3_0     conda-forge
arviz-base                 1.0.0            pyhd8ed1ab_0     conda-forge
arviz-plots                1.0.0            pyhc364b38_0     conda-forge
arviz-stats                1.0.0            pyh5442c79_0     conda-forge
arviz-stats-core           1.0.0            pyhc364b38_0     conda-forge
xarray                     2026.2.0         pyhcf101f3_0     conda-forge
xarray-einstats            0.10.0           pyhd8ed1ab_0     conda-forge

MWE:


import bambi as bmb
import pandas as pd

data = pd.DataFrame(
    {
        "driver": ["ham", "ham", "lec", "lec"],
        "pace": [100.3, 100.5, 100.1, 100.4],
        "session_type": ["R", "Q", "R", "Q"],
    }
)

model = bmb.Model(
    "pace ~ (0 + session_type|driver) + (1|driver)",
    data,
)
results = model.fit()

# Now try to predict
new_data = pd.DataFrame(
    {
        "driver": ["ham", "lec"],
        "session_type": ["R", "R"],
    }
)

predictions = model.predict(results, data=new_data, kind="response")
predictions

Stacktrace

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], [line 38](vscode-notebook-cell:?execution_count=1&line=38)
     30 # This should work but fails
     31 new_data = pd.DataFrame(
     32     {
     33         "driver": ["ham", "lec"],
     34         "session_type": ["R", "R"],
     35     }
     36 )
---> [38](vscode-notebook-cell:?execution_count=1&line=38) predictions = model.predict(results, data=new_data, kind="response")
     39 print(predictions)

File ~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/models.py:859, in Model.predict(self, idata, kind, data, inplace, include_group_specific, sample_new_groups, random_seed)
    856     idata = deepcopy(idata)
    858 # Populate the posterior in the InferenceData object with the likelihood parameters
--> [859](https://file+.vscode-resource.vscode-cdn.net/Users/nce8/github/bayesianf1/~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/models.py:859) idata = self._compute_likelihood_params(
    860     idata, data, include_group_specific, sample_new_groups
    861 )
    863 # Only if requested predict the predictive distribution
    864 if kind == "response":

File ~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/models.py:1013, in Model._compute_likelihood_params(self, idata, data, include_group_specific, sample_new_groups)
   1011 for name, component in self.distributional_components.items():
   1012     var_name = component.alias if component.alias else name
-> [1013](https://file+.vscode-resource.vscode-cdn.net/Users/nce8/github/bayesianf1/~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/models.py:1013)     means_dict[var_name] = component.predict(
   1014         idata, data, include_group_specific, hsgp_dict, sample_new_groups
   1015     )
   1017     # Drop var/dim if already present. Needed for out-of-sample predictions.
   1018     if var_name in idata.posterior.data_vars:

File ~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/model_components.py:160, in DistributionalComponent.predict(self, idata, data, include_group_specific, hsgp_dict, sample_new_groups)
    155     linear_predictor += self.predict_common(
    156         posterior, data, in_sample, to_stack_dims, design_matrix_dims, hsgp_dict
    157     )
    159 if self.design.group and include_group_specific:
--> [160](https://file+.vscode-resource.vscode-cdn.net/Users/nce8/github/bayesianf1/~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/model_components.py:160)     linear_predictor += self.predict_group_specific(
    161         posterior, data, in_sample, to_stack_dims, design_matrix_dims, sample_new_groups
    162     )
    164 # Sort dimensions
    165 linear_predictor = linear_predictor.transpose(*linear_predictor_dims)

File ~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/model_components.py:353, in DistributionalComponent.predict_group_specific(self, posterior, data, in_sample, to_stack_dims, design_matrix_dims, sample_new_groups)
    351 u = np.concatenate(u_arrays, axis=-1)
    352 u = xr.DataArray(u, dims=u_dims)
--> [353](https://file+.vscode-resource.vscode-cdn.net/Users/nce8/github/bayesianf1/~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/bambi/model_components.py:353) Z = xr.DataArray(Z, dims=design_matrix_dims)
    354 return xr.dot(Z, u)

File ~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/xarray/core/dataarray.py:462, in DataArray.__init__(self, data, coords, dims, name, attrs, indexes, fastpath)
    460 data = _check_data_shape(data, coords, dims)
    461 data = as_compatible_data(data)
--> [462](https://file+.vscode-resource.vscode-cdn.net/Users/nce8/github/bayesianf1/~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/xarray/core/dataarray.py:462) coords, dims = _infer_coords_and_dims(data.shape, coords, dims)
    463 variable = Variable(dims, data, attrs, fastpath=True)
    465 if not isinstance(coords, Coordinates):

File ~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/xarray/core/dataarray.py:167, in _infer_coords_and_dims(shape, coords, dims)
    165 dims_tuple = tuple(dims)
    166 if len(dims_tuple) != len(shape):
--> [167](https://file+.vscode-resource.vscode-cdn.net/Users/nce8/github/bayesianf1/~/miniforge3/envs/bayesianf1/lib/python3.14/site-packages/xarray/core/dataarray.py:167)     raise ValueError(
    168         "different number of dimensions on data "
    169         f"and dims: {len(shape)} vs {len(dims_tuple)}"
    170     )
    171 for d in dims_tuple:
    172     if not hashable(d):

ValueError: different number of dimensions on data and dims: 0 vs 2
1 Like

cc: @tcapretto

1 Like

I was running into the same error with bambi 0.16.0 but was able to fix by pinning to formulae 0.5.4.

You could probably also upgrade to the latest bambi version (0.17). But seems like something with this formulae release might have broken the predict method when include_group_specific = True for earlier bambi versions?

1 Like

I’m not crazy! Thanks @bmele . Not sure why I’m on 0.16 — this is from a pretty recent conda install. Oh well! Thanks for the help.

1 Like

Thanks @bmele for the solution!

1 Like