Can't seem to find the `variable_importance` attribute in the inference data.

[DevenVGokhale]

Hi team,
I'm working through an example using pymc-bart, adapted from one from the chapter 9 of the BAP text book. I'm attempting to extract and plot variable importance after fitting a BART model. While sampling works fine, the plot_variable_importance() function raises an error.

I'm unsure if I'm using the function incorrectly or if variable importance needs to be enabled/configured explicitly during model definition. I’d really appreciate your guidance on this!

📦 Setup
pymc version: 5.22
pymc-bart version: 0.9.2
Python: 3.12.10

My data has variables similar to this test data:

import polars as pl

df = pl.DataFrame({
    "cement": [100, 150, 120, 130, 160, 140, 110, 180],
    "slag": [0, 20, 10, 30, 0, 25, 5, 35],
    "ash": [10, 15, 20, 25, 10, 15, 20, 25],
    "water": [200, 180, 190, 210, 195, 185, 200, 205],
    "superplasticizer": [5, 6, 7, 8, 4, 5, 6, 7],
    "compressive_strength": [40, 45, 50, 42, 47, 43, 41, 46]
})

X = df.drop("compressive_strength").to_pandas()
y = df["compressive_strength"].to_pandas().to_numpy().ravel()

📈 Model Definition & Sampling

import pymc as pm
import pymc_bart as pmb

with pm.Model() as model_compress:
    μ = pmb.BART("μ", X, y)
    σ = pm.HalfNormal("σ", 1)
    y_obs = pm.Normal("y", mu=μ, sigma=σ, observed=y)

    idata_compress = pm.sample()

🧪 Problematic Call
pmb.plot_variable_importance(idata_compress, model_compress["μ"], X)

This raises the following error:
AxisError: axis -1 is out of bounds for array of dimension 0

🔍 What I’ve Tried
I’ve confirmed that model_compress["μ"] returns a valid BART object.
The posterior output looks like this:

idata_compress.posterior["μ"].shape  # (chains=4, draws=1000, μ_dim_0=772)

However, this fails:

idata_compress.posterior["μ"].attrs["variable_importance"]
# → raises KeyError or returns nothing

I will appreciate your feedback. Am I passing the wrong object to the plotting function?

Thank you in advance.

[aloctavodia]

You are using a new version of PyMC-BART, the API has changed. In the latest versions, the computation of variable importance and the plotting have been split into two functions.

What if you try this?

vi_results = pmb.compute_variable_importance(idata_compress, model_compress["μ"], X)
pmb.plot_variable_importance(vi_results);