Fix SHAP interaction output shape

intel · Sep 27, 2023 · 336cb94 · 336cb94
1 parent 598adb2
commit 336cb94
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/daal4py/mb/model_builders.py b/daal4py/mb/model_builders.py
@@ -222,7 +222,9 @@ def _predict_regression(
  )
  predict_result = predict_algo.compute(X, self.daal_model_)
 
- if pred_interactions:
+ if pred_contribs:
+ return predict_result.prediction.ravel().reshape((-1, X.shape[1] + 1))
+ elif pred_interactions:
  return predict_result.prediction.ravel().reshape(
  (-1, X.shape[1] + 1, X.shape[1] + 1)
  )