                          static_cast<float>(tree.thresholds[pos]), false,
                          left_zero_fraction, 0.0f);
 
-  RecurseTree(left_child, tree, tmp_path, paths, path_idx);
+  RecurseTree(left_child, tree, tmp_path, paths, path_idx, group, num_outputs);
 
   // Add left split to the path
   tmp_path->back() = gpu_treeshap::PathElement(
       0, tree.features[pos], 0, static_cast<float>(tree.thresholds[pos]), inf,
       false, 1.0 - left_zero_fraction, 0.0f);
 
-  RecurseTree(tree.children_right[pos], tree, tmp_path, paths, path_idx);
+  RecurseTree(tree.children_right[pos], tree, tmp_path, paths, path_idx, group,
+              num_outputs);
 
   tmp_path->pop_back();
 }


  // dispatch to the correct algorithm handler
  switch (feature_dependence) {
    case FEATURE_DEPENDENCE::independent:
      std::cerr << "FEATURE_DEPENDENCE::independent not yet supported\n";
      return;

    case FEATURE_DEPENDENCE::tree_path_dependent:
      if (interactions) {
        dense_tree_path_dependent_interactions_gpu(trees, data, out_contribs,
                                                   transform);
      } else {
        dense_tree_path_dependent_gpu(trees, data, out_contribs, transform);
      }
      return;

    case FEATURE_DEPENDENCE::global_path_dependent:
      std::cerr << "FEATURE_DEPENDENCE::global_path_dependent not supported\n";
      return;
    default:
      std::cerr << "Unknown feature dependence option\n";
      return;
  }
}

      device_data.GetDeviceAccessor();

  thrust::device_vector<float> phis((X.NumCols() + 1) * (X.NumCols() + 1) *
                                    X.NumRows() * trees.num_outputs);
  gpu_treeshap::GPUTreeShapInteractions(X, paths.begin(), paths.end(),
                                        trees.num_outputs, phis.data().get(),
                                        phis.size());
  // Add the base offset term to bias
  thrust::device_vector<double> base_offset(
      trees.base_offset, trees.base_offset + trees.num_outputs);
  auto counting = thrust::make_counting_iterator(size_t(0));
  auto d_phis = phis.data().get();
  auto d_base_offset = base_offset.data().get();
  size_t num_groups = trees.num_outputs;
  thrust::for_each(counting, counting + X.NumRows() * num_groups,
                   [=] __device__(size_t idx) {
                     size_t row_idx = idx / num_groups;
                     size_t group = idx % num_groups;
                     auto phi_idx = gpu_treeshap::IndexPhiInteractions(
                         row_idx, num_groups, group, X.NumCols(), X.NumCols(),
                         X.NumCols());
                     d_phis[phi_idx] += d_base_offset[group];
                   });
  // Shap uses a slightly different layout for multiclass
  thrust::device_vector<float> transposed_phis(phis.size());
  auto d_transposed_phis = transposed_phis.data();
  thrust::for_each(
      counting, counting + phis.size(), [=] __device__(size_t idx) {
        size_t old_shape[] = {X.NumRows(), num_groups, (X.NumCols() + 1),
                              (X.NumCols() + 1)};
        size_t old_idx[std::size(old_shape)];
        flat_idx_to_tensor_idx(idx, old_shape, old_idx);
        // Define new tensor format, switch num_groups axis to end
        size_t new_shape[] = {X.NumRows(), (X.NumCols() + 1), (X.NumCols() + 1),
                              num_groups};
        size_t new_idx[] = {old_idx[0], old_idx[2], old_idx[3], old_idx[1]};
        size_t transposed_idx = tensor_idx_to_flat_idx(new_shape, new_idx);
        d_transposed_phis[transposed_idx] = d_phis[idx];
      });
  thrust::copy(transposed_phis.begin(), transposed_phis.end(), out_contribs);
}

void dense_tree_shap_gpu(const TreeEnsemble &trees,
