   DeviceExplanationDataset::DenseDatasetWrapper X =
       device_data.GetDeviceAccessor();
 
-  thrust::device_vector<float> phis((X.NumCols() + 1) * X.NumRows());
-  gpu_treeshap::GPUTreeShap(X, paths.begin(), paths.end(), 1, phis.data().get(),
-                            phis.size());
+  thrust::device_vector<float> phis((X.NumCols() + 1) * X.NumRows() *
+                                    trees.num_outputs);
+  gpu_treeshap::GPUTreeShap(X, paths.begin(), paths.end(), trees.num_outputs,
+                            phis.data().get(), phis.size());
   // Add the base offset term to bias
-  auto base_offset = trees.base_offset[0];
+  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();
-  thrust::for_each(counting, counting + X.NumRows(),
-                   [=] __device__(size_t row_idx) {
+  auto d_base_offset = base_offset.data().get();
+  size_t num_groups = trees.num_outputs;
+  thrust::for_each(counting, counting + X.NumRows() * trees.num_outputs,
+                   [=] __device__(size_t idx) {
+                     size_t row_idx = idx / num_groups;
+                     size_t group = idx % num_groups;
                      auto phi_idx = gpu_treeshap::IndexPhi(
-                         row_idx, 1, 0, X.NumCols(), X.NumCols());
-                     d_phis[phi_idx] += base_offset;
+                         row_idx, num_groups, group, X.NumCols(), X.NumCols());
+                     d_phis[phi_idx] += d_base_offset[group];
                    });
-  thrust::copy(phis.begin(), phis.end(), out_contribs);
+
+  // 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)};
+        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), num_groups};
+        size_t new_idx[] = {old_idx[0], old_idx[2], 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);
 }
 
 inline void dense_tree_path_dependent_interactions_gpu(


  // 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;
  }
}

  size_t num_features;
  size_t num_rows;

public:
  DeviceExplanationDataset(const ExplanationDataset &host_data)
      : data(host_data.X, host_data.X + host_data.num_X * host_data.M),
        missing(host_data.X_missing,
                host_data.X_missing + host_data.num_X * host_data.M),
        num_features(host_data.M), num_rows(host_data.num_X) {}

  class DenseDatasetWrapper {
    const float *data;
