Shap value plotting error on Databricks but works locally

Question

I want to do a simple shap analysis and plot a shap.force_plot. I noticed that it works without any issues locally in a .ipynb file, but fails on Databricks with the following error message:

Visualization omitted, Javascript library not loaded!
Have you run `initjs()` in this notebook? If this notebook was from another user you must 
also trust this notebook (File -> Trust notebook). If you are viewing  this notebook on 
github the Javascript has been stripped for security. If you are using JupyterLab this 
error is because a JupyterLab extension has not yet been written.

Code:

import xgboost
import shap 

shap.initjs()

X, y = shap.datasets.boston()
bst = xgboost.train({"learning_rate": 0.01}, xgboost.DMatrix(X, label=y), 100)

explainer = shap.TreeExplainer(bst)
shap_values = explainer.shap_values(X)

shap.force_plot(explainer.expected_value, shap_values[0,:], X.iloc[0,:])

Is there any way to get the image plotting work on Databricks?

Answer 1

Let's try slightly different (matplotlib=True):

import xgboost
import shap 

X, y = shap.datasets.boston()
bst = xgboost.train({"learning_rate": 0.01}, xgboost.DMatrix(X, label=y), 100)

explainer = shap.TreeExplainer(bst)
shap_values = explainer.shap_values(X)

shap.force_plot(
    explainer.expected_value, 
    shap_values[0,:], 
    X.iloc[0,:], 
    matplotlib=True # <--
)