Use output-last SHAP C API shape

Author: RAMitchell

include/xgboost/c_api.h

@@ -1243,9 +1243,6 @@ XGB_DLL int XGBoosterPredictFromDMatrix(BoosterHandle handle, DMatrixHandle dmat
  *      Beginning iteration.
  *    "iteration_end": int
  *      End iteration.  Set to 0 to use all trees.
- *    "strict_shape": bool
- *      Whether output shapes should include the output-group dimension even when
- *      there is only one output group.
  *
  * @param out_values_shape Shape of feature SHAP values (copy before use).
  * @param out_values_dim Dimension of feature SHAP values.

python-package/xgboost/interpret.py

@@ -76,7 +76,6 @@ def _capi_shap_values(
         "algorithm": "auto",
         "iteration_begin": int(iteration_range[0]),
         "iteration_end": int(iteration_range[1]),
-        "strict_shape": False,
     }
     _check_call(
         _LIB.XGBoosterInterpretShapValues(
@@ -92,10 +91,12 @@ def _capi_shap_values(
             ctypes.byref(bias),
         )
     )
-    return (
-        _prediction_output(values_shape, values_dim, values, False),
-        _prediction_output(bias_shape, bias_dim, bias, False),
-    )
+    values_out = _prediction_output(values_shape, values_dim, values, False)
+    bias_out = _prediction_output(bias_shape, bias_dim, bias, False)
+    if values_out.shape[-1] == 1:
+        values_out = values_out[..., 0]
+        bias_out = bias_out[..., 0]
+    return values_out, bias_out
 
 
 def shap_values(  # pylint: disable=too-many-arguments

src/c_api/c_api.cc

@@ -1389,7 +1389,6 @@ XGB_DLL int XGBoosterInterpretShapValues(
   auto p_m = *static_cast<std::shared_ptr<DMatrix> *>(dmat);
   auto iteration_begin = OptionalArg<Integer>(config, "iteration_begin", Integer::Int{0});
   auto iteration_end = OptionalArg<Integer>(config, "iteration_end", Integer::Int{0});
-  bool strict_shape = OptionalArg<Boolean>(config, "strict_shape", false);
 
   learner->Predict(p_m, false, &entry.predictions, iteration_begin, iteration_end, false, false,
                    true, false, false);
@@ -1407,29 +1406,23 @@ XGB_DLL int XGBoosterInterpretShapValues(
   for (std::size_t row = 0; row < rows; ++row) {
     for (std::size_t group = 0; group < groups; ++group) {
       std::size_t contrib_offset = row * groups * (cols + 1) + group * (cols + 1);
-      std::size_t value_offset = row * groups * cols + group * cols;
-      std::copy_n(contribs.cbegin() + contrib_offset, cols, values.begin() + value_offset);
+      for (std::size_t col = 0; col < cols; ++col) {
+        std::size_t value_offset = row * cols * groups + col * groups + group;
+        values[value_offset] = contribs[contrib_offset + col];
+      }
       bias[row * groups + group] = contribs[contrib_offset + cols];
     }
   }
 
   auto &values_shape = local.prediction_shape;
   auto &bias_shape = local.prediction_shape_1;
-  if (groups == 1 && !strict_shape) {
-    values_shape.resize(2);
-    values_shape[0] = rows;
-    values_shape[1] = cols;
-    bias_shape.resize(1);
-    bias_shape[0] = rows;
-  } else {
-    values_shape.resize(3);
-    values_shape[0] = rows;
-    values_shape[1] = groups;
-    values_shape[2] = cols;
-    bias_shape.resize(2);
-    bias_shape[0] = rows;
-    bias_shape[1] = groups;
-  }
+  values_shape.resize(3);
+  values_shape[0] = rows;
+  values_shape[1] = cols;
+  values_shape[2] = groups;
+  bias_shape.resize(2);
+  bias_shape[0] = rows;
+  bias_shape[1] = groups;
 
   xgboost_CHECK_C_ARG_PTR(out_values_dim);
   xgboost_CHECK_C_ARG_PTR(out_values_shape);

tests/cpp/c_api/test_c_api.cc

@@ -742,7 +742,6 @@ TEST(CAPI, InterpretShapValues) {
   shap_config["algorithm"] = String{"auto"};
   shap_config["iteration_begin"] = Integer{0};
   shap_config["iteration_end"] = Integer{0};
-  shap_config["strict_shape"] = Boolean{false};
   auto sshap_config = Json::Dump(shap_config);
 
   bst_ulong const *values_shape{nullptr};
@@ -755,11 +754,13 @@ TEST(CAPI, InterpretShapValues) {
                                          &values_shape, &values_dim, &values, &bias_shape,
                                          &bias_dim, &bias),
             0);
-  ASSERT_EQ(values_dim, 2);
+  ASSERT_EQ(values_dim, 3);
   ASSERT_EQ(values_shape[0], n_samples);
   ASSERT_EQ(values_shape[1], n_features);
-  ASSERT_EQ(bias_dim, 1);
+  ASSERT_EQ(values_shape[2], 1);
+  ASSERT_EQ(bias_dim, 2);
   ASSERT_EQ(bias_shape[0], n_samples);
+  ASSERT_EQ(bias_shape[1], 1);
 
   Json pred_config{Object{}};
   pred_config["type"] = Integer{2};