From ee601c79eed55e573543d0f03adc649a84b98c33 Mon Sep 17 00:00:00 2001
From: John Walz <john@validmind.ai>
Date: Wed, 8 Jul 2026 12:02:29 -0400
Subject: [PATCH 1/3] Fix WeakspotsDiagnosis crash on multiclass targets and
 non-{0,1} binary labels

WeakspotsDiagnosis scores precision/recall/F1 per feature slice using
scikit-learn's defaults (average="binary", pos_label=1). This raises when a
slice does not look like a standard {0, 1} binary problem:

- Multiclass targets -> "Target is multiclass but average='binary'".
- Binary targets encoded outside {0, 1} -> "pos_label=1 is not a valid label".

The customer-reported case (ZD-704) is a genuinely multiclass target where one
feature bin happened to contain only two classes (e.g. {0, 4}); the old
per-slice scoring treated that slice as binary and crashed with
"pos_label=1 is not a valid label. It should be one of [0, 4]".

Resolve the averaging mode once from the full label set (union of true and
predicted labels across every dataset) and apply it to every slice:
multiclass -> average="macro"; binary -> average="binary" with the largest
class value as pos_label. For the standard {0, 1} case pos_label stays 1, so
existing results are unchanged. accuracy and any custom metric without an
`average` parameter are left untouched.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../sklearn/test_WeakspotsDiagnosis.py        | 195 ++++++++++++++++++
 .../sklearn/WeakspotsDiagnosis.py             |  63 +++++-
 2 files changed, 257 insertions(+), 1 deletion(-)

diff --git a/tests/unit_tests/model_validation/sklearn/test_WeakspotsDiagnosis.py b/tests/unit_tests/model_validation/sklearn/test_WeakspotsDiagnosis.py
index ff29cc113..92be8d2f4 100644
--- a/tests/unit_tests/model_validation/sklearn/test_WeakspotsDiagnosis.py
+++ b/tests/unit_tests/model_validation/sklearn/test_WeakspotsDiagnosis.py
@@ -1,10 +1,65 @@
 import unittest
 
+import numpy as np
+import pandas as pd
+from sklearn import metrics as skm
+from sklearn.linear_model import LogisticRegression
+
+import validmind as vm
 from validmind.tests.model_validation.sklearn.WeakspotsDiagnosis import (
+    WeakspotsDiagnosis,
+    _apply_averaging,
     _prepare_metrics_and_thresholds,
+    _resolve_averaging,
 )
 
 
+def _dataset_with_predictions(input_id, y_true, y_pred, model=None):
+    """Build a VMDataset with two numeric features and injected predictions.
+
+    Predictions are supplied via ``prediction_values`` so the model's ``predict``
+    is never called and the true/predicted label sets can be controlled exactly
+    (including binary targets encoded outside ``{0, 1}``). Passing an existing
+    ``model`` lets a train/test pair share one VMModel input.
+    """
+    n = len(y_true)
+    rng = np.random.RandomState(abs(hash(input_id)) % (2**32))
+    df = pd.DataFrame(
+        {
+            "f1": np.linspace(-2.0, 2.0, n),
+            "f2": rng.randn(n),
+            "target": y_true,
+        }
+    )
+    dataset = vm.init_dataset(
+        input_id=input_id, dataset=df, target_column="target", __log=False
+    )
+    if model is None:
+        estimator = LogisticRegression(max_iter=1000)
+        estimator.fit(df[["f1", "f2"]].to_numpy(), np.array(y_true))
+        model = vm.init_model(
+            input_id=f"{input_id}_model", model=estimator, __log=False
+        )
+    dataset.assign_predictions(model=model, prediction_values=y_pred)
+    return dataset, model
+
+
+def _train_test_pair(name, labels, seed=7, n=48):
+    rng = np.random.RandomState(seed)
+    train, model = _dataset_with_predictions(
+        f"{name}_train",
+        rng.choice(labels, size=n).tolist(),
+        rng.choice(labels, size=n).tolist(),
+    )
+    test, _ = _dataset_with_predictions(
+        f"{name}_test",
+        rng.choice(labels, size=n).tolist(),
+        rng.choice(labels, size=n).tolist(),
+        model=model,
+    )
+    return train, test, model
+
+
 class TestWeakspotsDiagnosisThresholds(unittest.TestCase):
     def test_partial_thresholds_use_defaults_for_plotting(self):
         _, plot_thresholds, pass_thresholds = _prepare_metrics_and_thresholds(
@@ -27,5 +82,145 @@ def test_partial_thresholds_subset_for_pass_fail(self):
         self.assertEqual(set(pass_thresholds.keys()), {"Accuracy", "F1"})
 
 
+class TestResolveAveraging(unittest.TestCase):
+    def test_binary_non_standard_encoding_selects_positive_label(self):
+        # ZD-704 follow-up: binary target encoded as {0, 4}. sklearn's default
+        # pos_label=1 is not a valid label, so the positive label must be resolved
+        # to the largest class value (4).
+        train, test, model = _train_test_pair("resolve_04", [0, 4])
+        self.assertEqual(_resolve_averaging([train, test], model), ("binary", 4))
+
+    def test_standard_binary_keeps_pos_label_one(self):
+        # {0, 1} must resolve to sklearn's own default so existing results are
+        # unchanged.
+        train, test, model = _train_test_pair("resolve_01", [0, 1])
+        self.assertEqual(_resolve_averaging([train, test], model), ("binary", 1))
+
+    def test_multiclass_selects_macro(self):
+        train, test, model = _train_test_pair("resolve_mc", [0, 1, 2])
+        self.assertEqual(_resolve_averaging([train, test], model), ("macro", None))
+
+    def test_multiclass_detected_from_predictions(self):
+        # True labels are binary but the model predicts a third class; the union of
+        # true and predicted labels must still be treated as multiclass.
+        train, model = _dataset_with_predictions(
+            "resolve_pred_mc",
+            [0, 1, 1, 0, 1, 0, 1, 0],
+            [0, 1, 2, 0, 1, 0, 2, 0],  # class 2 never appears in y_true
+        )
+        self.assertEqual(_resolve_averaging([train], model), ("macro", None))
+
+
+class TestApplyAveraging(unittest.TestCase):
+    def test_binary_binds_average_and_pos_label(self):
+        metrics, _, _ = _prepare_metrics_and_thresholds(metrics=None, thresholds=None)
+        prepared = _apply_averaging(metrics, "binary", 4)
+
+        y_true = np.array([0, 4, 0, 4])
+        y_pred = np.array([0, 4, 4, 4])
+        # accuracy has no `average`/`pos_label`, so it is left untouched.
+        self.assertIs(prepared["Accuracy"], skm.accuracy_score)
+        self.assertEqual(
+            prepared["Precision"](y_true, y_pred),
+            skm.precision_score(y_true, y_pred, average="binary", pos_label=4),
+        )
+
+    def test_macro_does_not_bind_pos_label(self):
+        metrics, _, _ = _prepare_metrics_and_thresholds(metrics=None, thresholds=None)
+        prepared = _apply_averaging(metrics, "macro", None)
+
+        y_true = np.array([0, 1, 2, 0, 1, 2])
+        y_pred = np.array([0, 1, 1, 0, 2, 2])
+        self.assertEqual(
+            prepared["F1"](y_true, y_pred),
+            skm.f1_score(y_true, y_pred, average="macro"),
+        )
+
+
+class TestWeakspotsDiagnosisEndToEnd(unittest.TestCase):
+    def test_binary_non_standard_encoding_runs(self):
+        # Regression test for ZD-704: WeakspotsDiagnosis previously raised
+        # "pos_label=1 is not a valid label. It should be one of [0, 4]".
+        train, test, model = _train_test_pair("e2e_04", [0, 4])
+        result = WeakspotsDiagnosis(datasets=[train, test], model=model)
+        self.assertGreater(len(result[0]), 0)
+        self.assertIsInstance(result[-1], bool)
+
+    def test_multiclass_runs(self):
+        # Previously raised "Target is multiclass but average='binary'".
+        train, test, model = _train_test_pair("e2e_mc", [0, 1, 2])
+        result = WeakspotsDiagnosis(datasets=[train, test], model=model)
+        self.assertGreater(len(result[0]), 0)
+
+    def test_multiclass_with_two_class_slice_runs(self):
+        # Faithful ZD-704 regression: the target is genuinely multiclass
+        # ({0, 1, 2, 3, 4}), but the lowest feature bin contains only two of those
+        # classes (0 and 4). The old per-slice scoring saw that slice as "binary"
+        # and raised "pos_label=1 is not a valid label. It should be one of [0, 4]".
+        # Resolving the averaging mode from the full (multiclass) label set makes
+        # every slice use macro averaging, so the sparse slice no longer fails.
+        n = 200
+        feature = np.linspace(0.0, 10.0, n)
+        rng = np.random.RandomState(0)
+
+        def build(input_id, model=None):
+            y = rng.choice([0, 1, 2, 3, 4], size=n)
+            y[:20] = rng.choice([0, 4], size=20)  # lowest bin has only {0, 4}
+            df = pd.DataFrame({"f1": feature, "f2": rng.randn(n), "target": y})
+            ds = vm.init_dataset(
+                input_id=input_id, dataset=df, target_column="target", __log=False
+            )
+            if model is None:
+                est = LogisticRegression(max_iter=1000)
+                est.fit(df[["f1", "f2"]].to_numpy(), y)
+                model = vm.init_model(
+                    input_id=f"{input_id}_model", model=est, __log=False
+                )
+            ds.assign_predictions(model=model)
+            return ds, model
+
+        train, model = build("e2e_mc_sparse_train")
+        test, _ = build("e2e_mc_sparse_test", model=model)
+
+        self.assertEqual(_resolve_averaging([train, test], model), ("macro", None))
+        result = WeakspotsDiagnosis(datasets=[train, test], model=model)
+        self.assertGreater(len(result[0]), 0)
+
+    def test_standard_binary_matches_default_precision(self):
+        # Backward compatibility: for {0, 1} the per-slice Precision values must be
+        # identical to scikit-learn's default binary precision so existing reports
+        # are unaffected by the fix.
+        train, test, model = _train_test_pair("e2e_01", [0, 1])
+        result = WeakspotsDiagnosis(datasets=[train, test], model=model)
+        df = result[0]
+
+        target = train.target_column
+        pred_col = train.prediction_column(model)
+        # Recompute the first non-empty slice's precision directly and compare.
+        checked = 0
+        for feature in train.feature_columns:
+            binned = train._df.copy()
+            binned["bin"] = pd.cut(binned[feature], bins=10)
+            for region, slice_df in binned.groupby("bin", observed=True):
+                if slice_df.empty:
+                    continue
+                y_true = slice_df[target].values
+                y_pred = slice_df[pred_col].astype(slice_df[target].dtype).values
+                expected = skm.precision_score(
+                    y_true, y_pred, average="binary", pos_label=1, zero_division=0
+                )
+                row = df[
+                    (df["Feature"] == feature)
+                    & (df["Slice"] == str(region))
+                    & (df["Dataset"] == train.input_id)
+                ]
+                self.assertAlmostEqual(row["Precision"].iloc[0], expected)
+                checked += 1
+                break
+            if checked:
+                break
+        self.assertGreater(checked, 0)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py b/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py
index c92fee871..1007dc9b4 100644
--- a/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py
+++ b/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py
@@ -2,9 +2,12 @@
 # Refer to the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-from typing import Callable, Dict, List, Optional, Tuple
+import functools
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import matplotlib.pyplot as plt
+import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
 import seaborn as sns
@@ -57,6 +60,58 @@ def _prepare_metrics_and_thresholds(
     return normalized_metrics, plot_thresholds, pass_thresholds
 
 
+def _resolve_averaging(
+    datasets: List[VMDataset], model: VMModel
+) -> Tuple[str, Optional[Any]]:
+    """Decide the averaging strategy for label-based metrics (precision, recall, F1).
+
+    scikit-learn's precision/recall/F1 default to ``average="binary"`` with
+    ``pos_label=1``. That raises for multiclass targets ("Target is multiclass but
+    average='binary'") and for binary targets encoded outside ``{0, 1}``
+    ("pos_label=1 is not a valid label"). The label space is taken from the union of
+    true and predicted labels across every dataset so that a slice whose labels
+    collapse to a subset is still scored consistently.
+
+    Returns an ``(average, pos_label)`` pair: ``("macro", None)`` for multiclass,
+    otherwise ``("binary", <positive label>)`` where the positive label is the
+    largest class value. This leaves the standard ``{0, 1}`` case on ``pos_label=1``
+    so existing binary results are unchanged.
+    """
+    labels = set()
+    for dataset in datasets:
+        labels.update(np.unique(dataset.y).tolist())
+        labels.update(np.unique(dataset.y_pred(model)).tolist())
+    labels = sorted(labels)
+
+    if len(labels) > 2:
+        return "macro", None
+    return "binary", labels[-1]
+
+
+def _apply_averaging(
+    metrics: Dict[str, Callable], average: str, pos_label: Optional[Any]
+) -> Dict[str, Callable]:
+    """Bind the resolved averaging options onto every metric that accepts them.
+
+    accuracy (and any custom metric without an ``average`` parameter) is left
+    untouched. ``pos_label`` is only bound for binary averaging, and
+    ``zero_division=0`` is bound when supported so that empty or single-class slices
+    score 0 instead of emitting scikit-learn warnings.
+    """
+    prepared = {}
+    for name, metric_fn in metrics.items():
+        params = inspect.signature(metric_fn).parameters
+        kwargs = {}
+        if "average" in params:
+            kwargs["average"] = average
+            if pos_label is not None and "pos_label" in params:
+                kwargs["pos_label"] = pos_label
+            if "zero_division" in params:
+                kwargs["zero_division"] = 0
+        prepared[name] = functools.partial(metric_fn, **kwargs) if kwargs else metric_fn
+    return prepared
+
+
 def _compute_metrics(
     results: dict,
     metrics: Dict[str, Callable],
@@ -264,6 +319,12 @@ def WeakspotsDiagnosis(
         metrics, thresholds
     )
 
+    # Bind averaging options so label-based metrics work for multiclass targets and
+    # for binary targets encoded outside {0, 1} (e.g. {0, 4}) instead of raising on
+    # scikit-learn's default average="binary"/pos_label=1.
+    average, pos_label = _resolve_averaging(datasets, model)
+    metrics = _apply_averaging(metrics, average, pos_label)
+
     results_headers = ["Slice", "Number of Records", "Feature"]
     results_headers.extend(metrics.keys())
 

From 0a80b55259b0f8252fd7681dd5ab2307f6f03d72 Mon Sep 17 00:00:00 2001
From: John Walz <john@validmind.ai>
Date: Wed, 8 Jul 2026 12:24:59 -0400
Subject: [PATCH 2/3] Fix OverfitDiagnosis and RobustnessDiagnosis multiclass
 crashes; share diagnosis metric helpers

OverfitDiagnosis and RobustnessDiagnosis have the same multiclass failures as
WeakspotsDiagnosis, plus an AUC-specific one:

- metric="f1"/"precision"/"recall" call scikit-learn with average="binary"
  (pos_label=1), raising "Target is multiclass but average='binary'" on
  multiclass targets and "pos_label=1 is not a valid label" on binary targets
  encoded outside {0, 1};
- the default metric "auc" calls roc_auc_score, which raises "multi_class must
  be in ('ovo', 'ovr')" for multiclass targets. A true multiclass ROC AUC is not
  possible here because the library retains only a single probability column, so
  fall back to the label-binarize convention already used by ClassifierPerformance
  (multiclass_roc_auc_score).

Extract the shared logic into _diagnosis_metrics.py:
- resolve_averaging/full_labels: decide (average, pos_label) once from the full
  label set (union of true and predicted labels) so every feature slice or
  perturbed copy is scored consistently;
- bind_averaging/apply_averaging: bind averaging onto metrics that accept it and
  pass accuracy/regression metrics through unchanged;
- multiclass_auc: robust per-subset label-binarize AUC (aligns columns on the
  full label set, scores only classes actually present, returns 0.0 otherwise).

WeakspotsDiagnosis now imports the shared helpers instead of its own copies. The
standard {0, 1} binary and regression paths are unchanged; the only behavior
change is that multiclass runs where it previously crashed. Adds unit tests for
the shared helpers and end-to-end tests for OverfitDiagnosis/RobustnessDiagnosis
across binary {0,1}, binary {0,4}, multiclass (default auc and f1), and regression.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../sklearn/test_OverfitDiagnosis.py          |  76 ++++++++++
 .../sklearn/test_RobustnessDiagnosis.py       | 108 ++++++++++++++
 .../sklearn/test_diagnosis_metrics.py         | 137 ++++++++++++++++++
 .../sklearn/OverfitDiagnosis.py               |  42 +++++-
 .../sklearn/RobustnessDiagnosis.py            |  46 +++++-
 .../sklearn/WeakspotsDiagnosis.py             |  60 +-------
 .../sklearn/_diagnosis_metrics.py             | 126 ++++++++++++++++
 7 files changed, 535 insertions(+), 60 deletions(-)
 create mode 100644 tests/unit_tests/model_validation/sklearn/test_OverfitDiagnosis.py
 create mode 100644 tests/unit_tests/model_validation/sklearn/test_RobustnessDiagnosis.py
 create mode 100644 tests/unit_tests/model_validation/sklearn/test_diagnosis_metrics.py
 create mode 100644 validmind/tests/model_validation/sklearn/_diagnosis_metrics.py

diff --git a/tests/unit_tests/model_validation/sklearn/test_OverfitDiagnosis.py b/tests/unit_tests/model_validation/sklearn/test_OverfitDiagnosis.py
new file mode 100644
index 000000000..b818bbd1b
--- /dev/null
+++ b/tests/unit_tests/model_validation/sklearn/test_OverfitDiagnosis.py
@@ -0,0 +1,76 @@
+import unittest
+
+import numpy as np
+import pandas as pd
+from sklearn.linear_model import LinearRegression, LogisticRegression
+
+import validmind as vm
+from validmind.tests.model_validation.sklearn.OverfitDiagnosis import OverfitDiagnosis
+
+
+def _pair(labels, seed=0, n=200, classification=True):
+    """Build a train/test VMDataset pair sharing one fitted VMModel."""
+
+    def build(input_id, s, model=None):
+        r = np.random.RandomState(s)
+        f1 = np.linspace(0.0, 10.0, n)
+        f2 = r.randn(n)
+        if classification:
+            y = r.choice(labels, size=n)
+            # ensure the lowest feature bin holds only a subset of the classes
+            y[:20] = r.choice(labels[:2], size=20)
+        else:
+            y = 3.0 * f1 + 2.0 * f2 + r.randn(n)
+        df = pd.DataFrame({"f1": f1, "f2": f2, "target": y})
+        ds = vm.init_dataset(
+            input_id=input_id, dataset=df, target_column="target", __log=False
+        )
+        if model is None:
+            est = (
+                LogisticRegression(max_iter=2000)
+                if classification
+                else LinearRegression()
+            )
+            est.fit(df[["f1", "f2"]].to_numpy(), y)
+            model = vm.init_model(input_id=f"{input_id}_m", model=est, __log=False)
+        ds.assign_predictions(model=model)
+        return ds, model
+
+    train, model = build(f"of_{seed}_train", seed)
+    test, _ = build(f"of_{seed}_test", seed + 1, model=model)
+    return train, test, model
+
+
+class TestOverfitDiagnosis(unittest.TestCase):
+    def test_multiclass_default_auc_runs(self):
+        # Regression test (ZD-704 family): default metric is "auc"; multiclass
+        # previously raised "multi_class must be in ('ovo', 'ovr')".
+        train, test, model = _pair([0, 1, 2, 3, 4], seed=1)
+        result = OverfitDiagnosis(model=model, datasets=[train, test])
+        self.assertIsNotNone(result)
+
+    def test_multiclass_f1_runs(self):
+        # metric="f1" previously raised "Target is multiclass but average='binary'".
+        train, test, model = _pair([0, 1, 2, 3, 4], seed=2)
+        result = OverfitDiagnosis(model=model, datasets=[train, test], metric="f1")
+        self.assertIsNotNone(result)
+
+    def test_binary_non_standard_labels_run(self):
+        # Binary target encoded as {0, 4}: default auc and f1 must both run.
+        train, test, model = _pair([0, 4], seed=3)
+        self.assertIsNotNone(OverfitDiagnosis(model=model, datasets=[train, test]))
+        self.assertIsNotNone(
+            OverfitDiagnosis(model=model, datasets=[train, test], metric="f1")
+        )
+
+    def test_standard_binary_runs(self):
+        train, test, model = _pair([0, 1], seed=4)
+        self.assertIsNotNone(OverfitDiagnosis(model=model, datasets=[train, test]))
+
+    def test_regression_runs(self):
+        train, test, model = _pair(None, seed=5, classification=False)
+        self.assertIsNotNone(OverfitDiagnosis(model=model, datasets=[train, test]))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit_tests/model_validation/sklearn/test_RobustnessDiagnosis.py b/tests/unit_tests/model_validation/sklearn/test_RobustnessDiagnosis.py
new file mode 100644
index 000000000..92f8300b6
--- /dev/null
+++ b/tests/unit_tests/model_validation/sklearn/test_RobustnessDiagnosis.py
@@ -0,0 +1,108 @@
+import unittest
+
+import numpy as np
+import pandas as pd
+from sklearn.linear_model import LinearRegression, LogisticRegression
+
+import validmind as vm
+from validmind.tests.model_validation.sklearn.RobustnessDiagnosis import (
+    RobustnessDiagnosis,
+)
+
+
+def _pair(labels, seed=0, n=200, classification=True):
+    """Build a train/test VMDataset pair sharing one fitted VMModel."""
+
+    def build(input_id, s, model=None):
+        r = np.random.RandomState(s)
+        f1 = np.linspace(0.0, 10.0, n)
+        f2 = r.randn(n)
+        if classification:
+            y = r.choice(labels, size=n)
+        else:
+            y = 3.0 * f1 + 2.0 * f2 + r.randn(n)
+        df = pd.DataFrame({"f1": f1, "f2": f2, "target": y})
+        ds = vm.init_dataset(
+            input_id=input_id, dataset=df, target_column="target", __log=False
+        )
+        if model is None:
+            est = (
+                LogisticRegression(max_iter=2000)
+                if classification
+                else LinearRegression()
+            )
+            est.fit(df[["f1", "f2"]].to_numpy(), y)
+            model = vm.init_model(input_id=f"{input_id}_m", model=est, __log=False)
+        ds.assign_predictions(model=model)
+        return ds, model
+
+    train, model = build(f"rb_{seed}_train", seed)
+    test, _ = build(f"rb_{seed}_test", seed + 1, model=model)
+    return train, test, model
+
+
+class TestRobustnessDiagnosis(unittest.TestCase):
+    # keep runs cheap: a single small perturbation is enough to exercise the paths
+    SMALL = [0.1]
+
+    def test_multiclass_default_auc_runs(self):
+        # Regression test (ZD-704 family): default metric is "auc"; multiclass
+        # previously raised "multi_class must be in ('ovo', 'ovr')".
+        train, test, model = _pair([0, 1, 2, 3, 4], seed=1)
+        result = RobustnessDiagnosis(
+            datasets=[train, test], model=model, scaling_factor_std_dev_list=self.SMALL
+        )
+        self.assertIsNotNone(result)
+
+    def test_multiclass_f1_runs(self):
+        # metric="f1" previously raised "Target is multiclass but average='binary'".
+        train, test, model = _pair([0, 1, 2, 3, 4], seed=2)
+        result = RobustnessDiagnosis(
+            datasets=[train, test],
+            model=model,
+            metric="f1",
+            scaling_factor_std_dev_list=self.SMALL,
+        )
+        self.assertIsNotNone(result)
+
+    def test_binary_non_standard_labels_run(self):
+        train, test, model = _pair([0, 4], seed=3)
+        self.assertIsNotNone(
+            RobustnessDiagnosis(
+                datasets=[train, test],
+                model=model,
+                scaling_factor_std_dev_list=self.SMALL,
+            )
+        )
+        self.assertIsNotNone(
+            RobustnessDiagnosis(
+                datasets=[train, test],
+                model=model,
+                metric="f1",
+                scaling_factor_std_dev_list=self.SMALL,
+            )
+        )
+
+    def test_standard_binary_runs(self):
+        train, test, model = _pair([0, 1], seed=4)
+        self.assertIsNotNone(
+            RobustnessDiagnosis(
+                datasets=[train, test],
+                model=model,
+                scaling_factor_std_dev_list=self.SMALL,
+            )
+        )
+
+    def test_regression_runs(self):
+        train, test, model = _pair(None, seed=5, classification=False)
+        self.assertIsNotNone(
+            RobustnessDiagnosis(
+                datasets=[train, test],
+                model=model,
+                scaling_factor_std_dev_list=self.SMALL,
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit_tests/model_validation/sklearn/test_diagnosis_metrics.py b/tests/unit_tests/model_validation/sklearn/test_diagnosis_metrics.py
new file mode 100644
index 000000000..d15581c5a
--- /dev/null
+++ b/tests/unit_tests/model_validation/sklearn/test_diagnosis_metrics.py
@@ -0,0 +1,137 @@
+import unittest
+
+import numpy as np
+import pandas as pd
+from sklearn import metrics as skm
+from sklearn.linear_model import LogisticRegression
+from sklearn.preprocessing import LabelBinarizer
+
+import validmind as vm
+from validmind.tests.model_validation.sklearn._diagnosis_metrics import (
+    apply_averaging,
+    bind_averaging,
+    full_labels,
+    multiclass_auc,
+    resolve_averaging,
+)
+
+
+def _dataset_with_predictions(input_id, y_true, y_pred, model=None):
+    """Build a VMDataset with injected predictions to control label sets exactly."""
+    n = len(y_true)
+    rng = np.random.RandomState(abs(hash(input_id)) % (2**32))
+    df = pd.DataFrame(
+        {
+            "f1": np.linspace(-2.0, 2.0, n),
+            "f2": rng.randn(n),
+            "target": y_true,
+        }
+    )
+    dataset = vm.init_dataset(
+        input_id=input_id, dataset=df, target_column="target", __log=False
+    )
+    if model is None:
+        estimator = LogisticRegression(max_iter=1000)
+        estimator.fit(df[["f1", "f2"]].to_numpy(), np.array(y_true))
+        model = vm.init_model(
+            input_id=f"{input_id}_model", model=estimator, __log=False
+        )
+    dataset.assign_predictions(model=model, prediction_values=y_pred)
+    return dataset, model
+
+
+class TestFullLabelsAndResolveAveraging(unittest.TestCase):
+    def test_full_labels_unions_true_and_predicted(self):
+        # class 4 appears only in predictions; it must still be part of the space.
+        ds, model = _dataset_with_predictions(
+            "fl", [0, 1, 2, 0, 1, 2], [0, 1, 2, 0, 4, 2]
+        )
+        self.assertEqual(full_labels([ds], model), [0, 1, 2, 4])
+
+    def test_resolve_binary_non_standard(self):
+        ds, model = _dataset_with_predictions("rb", [0, 4, 0, 4], [0, 4, 4, 0])
+        self.assertEqual(resolve_averaging([ds], model), ("binary", 4))
+
+    def test_resolve_standard_binary(self):
+        ds, model = _dataset_with_predictions("rs", [0, 1, 0, 1], [0, 1, 1, 0])
+        self.assertEqual(resolve_averaging([ds], model), ("binary", 1))
+
+    def test_resolve_multiclass(self):
+        ds, model = _dataset_with_predictions(
+            "rm", [0, 1, 2, 0, 1, 2], [0, 1, 2, 1, 1, 2]
+        )
+        self.assertEqual(resolve_averaging([ds], model), ("macro", None))
+
+
+class TestBindAveraging(unittest.TestCase):
+    def test_none_average_is_passthrough(self):
+        self.assertIs(bind_averaging(skm.f1_score, None, None), skm.f1_score)
+
+    def test_accuracy_passthrough(self):
+        # accuracy has no `average` parameter, so it must be returned unchanged.
+        self.assertIs(
+            bind_averaging(skm.accuracy_score, "macro", None), skm.accuracy_score
+        )
+
+    def test_binary_binds_pos_label(self):
+        fn = bind_averaging(skm.precision_score, "binary", 4)
+        y_true = np.array([0, 4, 0, 4])
+        y_pred = np.array([0, 4, 4, 4])
+        self.assertEqual(
+            fn(y_true, y_pred),
+            skm.precision_score(y_true, y_pred, average="binary", pos_label=4),
+        )
+
+    def test_macro_omits_pos_label(self):
+        fn = bind_averaging(skm.f1_score, "macro", None)
+        y_true = np.array([0, 1, 2, 0, 1, 2])
+        y_pred = np.array([0, 1, 1, 0, 2, 2])
+        self.assertEqual(
+            fn(y_true, y_pred), skm.f1_score(y_true, y_pred, average="macro")
+        )
+
+    def test_apply_averaging_maps_over_dict(self):
+        prepared = apply_averaging(
+            {"acc": skm.accuracy_score, "f1": skm.f1_score}, "macro", None
+        )
+        self.assertIs(prepared["acc"], skm.accuracy_score)
+        y_true = np.array([0, 1, 2, 0, 1, 2])
+        y_pred = np.array([0, 1, 1, 0, 2, 2])
+        self.assertEqual(
+            prepared["f1"](y_true, y_pred),
+            skm.f1_score(y_true, y_pred, average="macro"),
+        )
+
+
+class TestMulticlassAuc(unittest.TestCase):
+    LABELS = [0, 1, 2, 3, 4]
+
+    def test_matches_label_binarize_when_all_classes_present(self):
+        # When every class appears, the helper must equal the label-binarize AUC
+        # convention used by ClassifierPerformance.
+        y_true = np.array([0, 1, 2, 3, 4, 0, 1, 2, 3, 4])
+        y_pred = np.array([0, 1, 2, 2, 4, 0, 3, 2, 3, 4])
+        lb = LabelBinarizer().fit(y_true)
+        expected = skm.roc_auc_score(
+            lb.transform(y_true), lb.transform(y_pred), average="macro"
+        )
+        self.assertAlmostEqual(multiclass_auc(y_true, y_pred, self.LABELS), expected)
+
+    def test_subset_slice_does_not_raise(self):
+        # A slice with only two of the classes (0 and 4) plus a stray prediction.
+        y_true = np.array([0, 0, 4, 4, 0])
+        y_pred = np.array([0, 2, 4, 0, 4])
+        value = multiclass_auc(y_true, y_pred, self.LABELS)
+        self.assertIsInstance(value, float)
+        self.assertGreaterEqual(value, 0.0)
+        self.assertLessEqual(value, 1.0)
+
+    def test_single_class_slice_returns_zero(self):
+        # No class has both positive and negative examples -> nothing scorable.
+        y_true = np.array([4, 4, 4])
+        y_pred = np.array([4, 0, 4])
+        self.assertEqual(multiclass_auc(y_true, y_pred, self.LABELS), 0.0)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py b/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py
index ecff1006d..031f0cebf 100644
--- a/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py
+++ b/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py
@@ -15,6 +15,8 @@
 from validmind.logging import get_logger
 from validmind.vm_models import VMDataset, VMModel
 
+from ._diagnosis_metrics import bind_averaging, full_labels, multiclass_auc
+
 logger = get_logger(__name__)
 
 # TODO: A couple of improvements here could be to:
@@ -96,6 +98,10 @@ def _compute_metrics(
     feature_column: str,
     metric: str,
     is_classification: bool,
+    average: str = None,
+    pos_label=None,
+    labels: list = None,
+    is_multiclass: bool = False,
 ) -> None:
     results["slice"].append(str(region))
     results["shape"].append(df_region.shape[0])
@@ -115,11 +121,25 @@ def _compute_metrics(
         if len(np.unique(y_true)) == 1:
             return results[metric].append(0)
 
+        # The library retains only a single probability column, so a
+        # probability-based multiclass ROC AUC is not possible; fall back to the
+        # label-binarize convention used elsewhere for multiclass targets.
+        if is_multiclass:
+            return results[metric].append(
+                multiclass_auc(y_true, df_region[pred_column].values, labels)
+            )
+
         return results[metric].append(
             metric_func(y_true, df_region[prob_column].values)
         )
 
-    return results[metric].append(metric_func(y_true, df_region[pred_column].values))
+    # Bind averaging so precision/recall/F1 handle multiclass targets and binary
+    # targets encoded outside {0, 1}; accuracy and regression metrics pass through.
+    return results[metric].append(
+        bind_averaging(metric_func, average, pos_label)(
+            y_true, df_region[pred_column].values
+        )
+    )
 
 
 def _plot_overfit_regions(
@@ -253,6 +273,18 @@ def OverfitDiagnosis(
         train_df[prob_column] = datasets[0].y_prob(model)
         test_df[prob_column] = datasets[1].y_prob(model)
 
+    # Resolve the label space once so every feature slice is scored consistently:
+    # multiclass targets use macro averaging (and label-binarize AUC), binary
+    # targets keep the positive label so non-{0, 1} encodings don't break.
+    if is_classification:
+        labels = full_labels(datasets, model)
+        is_multiclass = len(labels) > 2
+        average, pos_label = (
+            ("macro", None) if is_multiclass else ("binary", labels[-1])
+        )
+    else:
+        labels, is_multiclass, average, pos_label = None, False, None, None
+
     test_results = []
     figures = []
     results_headers = ["slice", "shape", "feature", metric]
@@ -277,6 +309,10 @@ def OverfitDiagnosis(
                 pred_column=pred_column,
                 metric=metric,
                 is_classification=is_classification,
+                average=average,
+                pos_label=pos_label,
+                labels=labels,
+                is_multiclass=is_multiclass,
             )
             df_test_region = test_df[
                 (test_df[feature_column] > region.left)
@@ -292,6 +328,10 @@ def OverfitDiagnosis(
                 pred_column=pred_column,
                 metric=metric,
                 is_classification=is_classification,
+                average=average,
+                pos_label=pos_label,
+                labels=labels,
+                is_multiclass=is_multiclass,
             )
 
         results = _prepare_results(results_train, results_test, metric)
diff --git a/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py b/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
index b7639583d..5343ef274 100644
--- a/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
+++ b/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
@@ -17,6 +17,8 @@
 from validmind.logging import get_logger
 from validmind.vm_models import VMDataset, VMModel
 
+from ._diagnosis_metrics import bind_averaging, full_labels, multiclass_auc
+
 logger = get_logger(__name__)
 
 DEFAULT_DECAY_THRESHOLD = 0.05
@@ -86,7 +88,14 @@ def _add_noise_std_dev(
 
 
 def _compute_metric(
-    dataset: VMDataset, model: VMModel, X: pd.DataFrame, metric: str
+    dataset: VMDataset,
+    model: VMModel,
+    X: pd.DataFrame,
+    metric: str,
+    average: str = None,
+    pos_label=None,
+    labels: list = None,
+    is_multiclass: bool = False,
 ) -> float:
     if metric not in PERFORMANCE_METRICS:
         raise ValueError(
@@ -94,13 +103,23 @@ def _compute_metric(
         )
 
     if metric == "auc":
+        # Only a single probability column is available, so a probability-based
+        # multiclass ROC AUC is not possible; fall back to the label-binarize
+        # convention used elsewhere for multiclass targets.
+        if is_multiclass:
+            return multiclass_auc(dataset.y, model.predict(X), labels)
         try:
             y_proba = model.predict_proba(X)
         except MissingOrInvalidModelPredictFnError:
             y_proba = model.predict(X)
         return metrics.roc_auc_score(dataset.y, y_proba)
 
-    return PERFORMANCE_METRICS[metric]["function"](dataset.y, model.predict(X))
+    # Bind averaging so precision/recall/F1 handle multiclass targets and binary
+    # targets encoded outside {0, 1}; accuracy and regression metrics pass through.
+    metric_func = bind_averaging(
+        PERFORMANCE_METRICS[metric]["function"], average, pos_label
+    )
+    return metric_func(dataset.y, model.predict(X))
 
 
 def _compute_gap(result: dict, metric: str) -> float:
@@ -260,13 +279,26 @@ def RobustnessDiagnosis(
     - The test may not account for more complex or unstructured noise patterns that could affect model robustness.
     """
     # TODO: use single dataset
+    is_classification = bool(datasets[0].probability_column(model))
     if not metric:
         metric = (
             DEFAULT_CLASSIFICATION_METRIC
-            if datasets[0].probability_column(model)
+            if is_classification
             else DEFAULT_REGRESSION_METRIC
         )
 
+    # Resolve the label space once: multiclass targets use macro averaging (and
+    # label-binarize AUC), binary targets keep the positive label so non-{0, 1}
+    # encodings don't break scikit-learn's binary averaging.
+    if is_classification:
+        labels = full_labels(datasets, model)
+        is_multiclass = len(labels) > 2
+        average, pos_label = (
+            ("macro", None) if is_multiclass else ("binary", labels[-1])
+        )
+    else:
+        labels, is_multiclass, average, pos_label = None, False, None, None
+
     results = [{} for _ in range(len(datasets))]
 
     # add baseline results (no perturbation)
@@ -281,6 +313,10 @@ def RobustnessDiagnosis(
                 model=model,
                 X=dataset.x_df(),
                 metric=metric,
+                average=average,
+                pos_label=pos_label,
+                labels=labels,
+                is_multiclass=is_multiclass,
             )
         ]
         result["Performance Decay"] = [0.0]
@@ -307,6 +343,10 @@ def RobustnessDiagnosis(
                     model=model,
                     X=temp_df,
                     metric=metric,
+                    average=average,
+                    pos_label=pos_label,
+                    labels=labels,
+                    is_multiclass=is_multiclass,
                 )
             )
             result["Performance Decay"].append(_compute_gap(result, metric))
diff --git a/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py b/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py
index 1007dc9b4..bef73c9fe 100644
--- a/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py
+++ b/validmind/tests/model_validation/sklearn/WeakspotsDiagnosis.py
@@ -2,12 +2,9 @@
 # Refer to the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-import functools
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 import matplotlib.pyplot as plt
-import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
 import seaborn as sns
@@ -16,6 +13,9 @@
 from validmind import tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
+from ._diagnosis_metrics import apply_averaging as _apply_averaging
+from ._diagnosis_metrics import resolve_averaging as _resolve_averaging
+
 DEFAULT_METRICS = {
     "accuracy": metrics.accuracy_score,
     "precision": metrics.precision_score,
@@ -60,58 +60,6 @@ def _prepare_metrics_and_thresholds(
     return normalized_metrics, plot_thresholds, pass_thresholds
 
 
-def _resolve_averaging(
-    datasets: List[VMDataset], model: VMModel
-) -> Tuple[str, Optional[Any]]:
-    """Decide the averaging strategy for label-based metrics (precision, recall, F1).
-
-    scikit-learn's precision/recall/F1 default to ``average="binary"`` with
-    ``pos_label=1``. That raises for multiclass targets ("Target is multiclass but
-    average='binary'") and for binary targets encoded outside ``{0, 1}``
-    ("pos_label=1 is not a valid label"). The label space is taken from the union of
-    true and predicted labels across every dataset so that a slice whose labels
-    collapse to a subset is still scored consistently.
-
-    Returns an ``(average, pos_label)`` pair: ``("macro", None)`` for multiclass,
-    otherwise ``("binary", <positive label>)`` where the positive label is the
-    largest class value. This leaves the standard ``{0, 1}`` case on ``pos_label=1``
-    so existing binary results are unchanged.
-    """
-    labels = set()
-    for dataset in datasets:
-        labels.update(np.unique(dataset.y).tolist())
-        labels.update(np.unique(dataset.y_pred(model)).tolist())
-    labels = sorted(labels)
-
-    if len(labels) > 2:
-        return "macro", None
-    return "binary", labels[-1]
-
-
-def _apply_averaging(
-    metrics: Dict[str, Callable], average: str, pos_label: Optional[Any]
-) -> Dict[str, Callable]:
-    """Bind the resolved averaging options onto every metric that accepts them.
-
-    accuracy (and any custom metric without an ``average`` parameter) is left
-    untouched. ``pos_label`` is only bound for binary averaging, and
-    ``zero_division=0`` is bound when supported so that empty or single-class slices
-    score 0 instead of emitting scikit-learn warnings.
-    """
-    prepared = {}
-    for name, metric_fn in metrics.items():
-        params = inspect.signature(metric_fn).parameters
-        kwargs = {}
-        if "average" in params:
-            kwargs["average"] = average
-            if pos_label is not None and "pos_label" in params:
-                kwargs["pos_label"] = pos_label
-            if "zero_division" in params:
-                kwargs["zero_division"] = 0
-        prepared[name] = functools.partial(metric_fn, **kwargs) if kwargs else metric_fn
-    return prepared
-
-
 def _compute_metrics(
     results: dict,
     metrics: Dict[str, Callable],
diff --git a/validmind/tests/model_validation/sklearn/_diagnosis_metrics.py b/validmind/tests/model_validation/sklearn/_diagnosis_metrics.py
new file mode 100644
index 000000000..86a98184f
--- /dev/null
+++ b/validmind/tests/model_validation/sklearn/_diagnosis_metrics.py
@@ -0,0 +1,126 @@
+# Copyright © 2023-2026 ValidMind Inc. All rights reserved.
+# Refer to the LICENSE file in the root of this repository for details.
+# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
+
+"""Shared metric helpers for the per-region classification diagnosis tests.
+
+WeakspotsDiagnosis, OverfitDiagnosis and RobustnessDiagnosis all score
+classification metrics (precision/recall/F1/AUC) on subsets of the data -- feature
+bins or noise-perturbed copies. scikit-learn's defaults break on those subsets:
+
+- precision/recall/F1 default to ``average="binary"`` with ``pos_label=1``, which
+  raises for multiclass targets ("Target is multiclass but average='binary'") and
+  for binary targets encoded outside ``{0, 1}`` ("pos_label=1 is not a valid
+  label");
+- ``roc_auc_score`` defaults to ``multi_class="raise"`` and needs a per-class
+  probability matrix, which the library does not retain for multiclass models
+  (only a single probability column is stored).
+
+The averaging strategy is resolved once from the full label set so every subset is
+scored consistently, and multiclass AUC falls back to the label-binarize
+convention already used by ``ClassifierPerformance``.
+"""
+
+import functools
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+from sklearn.metrics import roc_auc_score
+from sklearn.preprocessing import LabelBinarizer
+
+from validmind.vm_models import VMDataset, VMModel
+
+
+def full_labels(datasets: List[VMDataset], model: VMModel) -> List[Any]:
+    """Return the sorted union of true and predicted labels across every dataset.
+
+    Predictions are included because scikit-learn derives a metric's target type
+    from the union of ``y_true`` and ``y_pred``: a subset whose true labels
+    collapse to two classes is still multiclass if the model predicts a third.
+    """
+    labels = set()
+    for dataset in datasets:
+        labels.update(np.unique(dataset.y).tolist())
+        labels.update(np.unique(dataset.y_pred(model)).tolist())
+    return sorted(labels)
+
+
+def resolve_averaging(
+    datasets: List[VMDataset], model: VMModel
+) -> Tuple[str, Optional[Any]]:
+    """Return the ``(average, pos_label)`` for precision/recall/F1.
+
+    ``("macro", None)`` for multiclass; ``("binary", <largest label>)`` otherwise,
+    which leaves the standard ``{0, 1}`` case on ``pos_label=1`` so existing binary
+    results are unchanged.
+    """
+    labels = full_labels(datasets, model)
+    if len(labels) > 2:
+        return "macro", None
+    return "binary", labels[-1]
+
+
+def bind_averaging(
+    metric_fn: Callable, average: Optional[str], pos_label: Optional[Any]
+) -> Callable:
+    """Bind averaging options onto a metric that accepts them; pass others through.
+
+    accuracy, regression metrics and any callable without an ``average`` parameter
+    are returned unchanged. ``pos_label`` is only bound for binary averaging, and
+    ``zero_division=0`` when supported so empty or single-class subsets score 0
+    instead of emitting scikit-learn warnings. ``average=None`` (regression) is a
+    no-op.
+    """
+    if average is None:
+        return metric_fn
+
+    params = inspect.signature(metric_fn).parameters
+    kwargs = {}
+    if "average" in params:
+        kwargs["average"] = average
+        if pos_label is not None and "pos_label" in params:
+            kwargs["pos_label"] = pos_label
+        if "zero_division" in params:
+            kwargs["zero_division"] = 0
+    return functools.partial(metric_fn, **kwargs) if kwargs else metric_fn
+
+
+def apply_averaging(
+    metrics: Dict[str, Callable], average: Optional[str], pos_label: Optional[Any]
+) -> Dict[str, Callable]:
+    """Bind averaging options onto every metric in a ``name -> callable`` mapping."""
+    return {
+        name: bind_averaging(metric_fn, average, pos_label)
+        for name, metric_fn in metrics.items()
+    }
+
+
+def multiclass_auc(y_true, y_pred, labels: List[Any]) -> float:
+    """Compute ROC AUC for multiclass targets from predicted labels.
+
+    The library retains only a single probability column, so a probability-based
+    multiclass ROC AUC is not possible; this mirrors ``ClassifierPerformance`` by
+    binarizing labels instead. Binarization is fit on the full label set so columns
+    stay aligned across subsets that omit some classes; only classes actually
+    present in ``y_true`` (with both positive and negative examples) are scored,
+    then macro-averaged. Returns ``0.0`` when no class is scorable.
+    """
+    binarizer = LabelBinarizer().fit(np.asarray(labels))
+    true_bin = binarizer.transform(np.asarray(y_true))
+    pred_bin = binarizer.transform(np.asarray(y_pred))
+
+    # LabelBinarizer collapses to a single column only for a two-class fit; the
+    # diagnosis tests call this for multiclass, but guard defensively.
+    if true_bin.shape[1] == 1:
+        true_bin = np.column_stack([1 - true_bin, true_bin])
+        pred_bin = np.column_stack([1 - pred_bin, pred_bin])
+
+    n_rows = len(y_true)
+    scorable = [
+        i for i in range(true_bin.shape[1]) if 0 < int(true_bin[:, i].sum()) < n_rows
+    ]
+    if not scorable:
+        return 0.0
+
+    return roc_auc_score(true_bin[:, scorable], pred_bin[:, scorable], average="macro")

From b192adb3f89b75a3d4b0931bb1847fa6f7a82e22 Mon Sep 17 00:00:00 2001
From: John Walz <john@validmind.ai>
Date: Wed, 8 Jul 2026 12:46:08 -0400
Subject: [PATCH 3/3] Fix CalibrationCurve and SHAPGlobalImportance multiclass
 failures

Two more sklearn tests fail on multiclass targets (found by sweeping every
classification test against a multiclass model):

- CalibrationCurve wraps sklearn's binary-only calibration_curve, which raised a
  cryptic "y_true takes value in {...} and pos_label is not specified" error on
  multiclass. Skip cleanly with SkipTestError instead, matching ROCCurve and
  PrecisionRecallCurve.
- SHAPGlobalImportance's select_shap_values only handled the older SHAP list
  layout (one 2D array per class). Newer SHAP returns a single
  (samples, features, classes) array for multiclass, which fell through to the
  regression branch and was passed whole to shap.summary_plot, raising
  "IndexError: index N is out of bounds for axis 1". Handle the 3D array by
  selecting along the trailing class axis; multiclass now works when
  class_of_interest is given and raises a clear, actionable error when it is not.

Binary and regression paths are unchanged. Adds unit tests for select_shap_values
(list, 3D-array, 2D-array, and error cases) and CalibrationCurve (multiclass skip,
binary runs).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../sklearn/test_CalibrationCurve.py          | 43 ++++++++++++++
 .../sklearn/test_SHAPGlobalImportance.py      | 52 ++++++++++++++++
 .../sklearn/CalibrationCurve.py               | 10 ++++
 .../sklearn/SHAPGlobalImportance.py           | 59 +++++++++++++------
 4 files changed, 145 insertions(+), 19 deletions(-)
 create mode 100644 tests/unit_tests/model_validation/sklearn/test_CalibrationCurve.py
 create mode 100644 tests/unit_tests/model_validation/sklearn/test_SHAPGlobalImportance.py

diff --git a/tests/unit_tests/model_validation/sklearn/test_CalibrationCurve.py b/tests/unit_tests/model_validation/sklearn/test_CalibrationCurve.py
new file mode 100644
index 000000000..cd5b33bc0
--- /dev/null
+++ b/tests/unit_tests/model_validation/sklearn/test_CalibrationCurve.py
@@ -0,0 +1,43 @@
+import unittest
+
+import numpy as np
+import pandas as pd
+from sklearn.linear_model import LogisticRegression
+
+import validmind as vm
+from validmind.errors import SkipTestError
+from validmind.tests.model_validation.sklearn.CalibrationCurve import CalibrationCurve
+
+
+def _dataset_and_model(labels, seed):
+    n = 200
+    rng = np.random.RandomState(seed)
+    y = rng.choice(labels, size=n)
+    df = pd.DataFrame(
+        {"f1": np.linspace(0.0, 10.0, n), "f2": rng.randn(n), "target": y}
+    )
+    ds = vm.init_dataset(
+        input_id=f"cc_{seed}", dataset=df, target_column="target", __log=False
+    )
+    est = LogisticRegression(max_iter=1000).fit(df[["f1", "f2"]].to_numpy(), y)
+    model = vm.init_model(input_id=f"cc_{seed}_m", model=est, __log=False)
+    ds.assign_predictions(model=model)
+    return ds, model
+
+
+class TestCalibrationCurve(unittest.TestCase):
+    def test_multiclass_skips_cleanly(self):
+        # sklearn's calibration_curve is binary-only and raised a cryptic
+        # "pos_label is not specified" error on multiclass; it must skip instead.
+        ds, model = _dataset_and_model([0, 1, 2, 3, 4], seed=1)
+        with self.assertRaises(SkipTestError):
+            CalibrationCurve(model=model, dataset=ds)
+
+    def test_binary_runs(self):
+        ds, model = _dataset_and_model([0, 1], seed=2)
+        result = CalibrationCurve(model=model, dataset=ds)
+        self.assertIsNotNone(result)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit_tests/model_validation/sklearn/test_SHAPGlobalImportance.py b/tests/unit_tests/model_validation/sklearn/test_SHAPGlobalImportance.py
new file mode 100644
index 000000000..fcc1ccfe4
--- /dev/null
+++ b/tests/unit_tests/model_validation/sklearn/test_SHAPGlobalImportance.py
@@ -0,0 +1,52 @@
+import unittest
+
+import numpy as np
+
+from validmind.tests.model_validation.sklearn.SHAPGlobalImportance import (
+    select_shap_values,
+)
+
+
+class TestSelectShapValues(unittest.TestCase):
+    def test_binary_or_regression_2d_array_passthrough(self):
+        # A single (samples, features) array (binary or regression) is used as-is.
+        values = np.arange(6.0).reshape(3, 2)
+        np.testing.assert_array_equal(select_shap_values(values), values)
+
+    def test_list_binary_defaults_to_class_one(self):
+        class0 = np.zeros((3, 2))
+        class1 = np.ones((3, 2))
+        np.testing.assert_array_equal(select_shap_values([class0, class1]), class1)
+
+    def test_list_multiclass_selects_requested_class(self):
+        classes = [np.full((3, 2), i, dtype=float) for i in range(5)]
+        np.testing.assert_array_equal(
+            select_shap_values(classes, class_of_interest=3),
+            np.full((3, 2), 3.0),
+        )
+
+    def test_3d_array_multiclass_selects_trailing_class_axis(self):
+        # Newer SHAP returns (samples, features, classes); regression test for the
+        # IndexError raised when this array was passed straight to summary_plot.
+        samples, features, classes = 4, 2, 5
+        values = np.arange(samples * features * classes, dtype=float).reshape(
+            samples, features, classes
+        )
+        selected = select_shap_values(values, class_of_interest=2)
+        self.assertEqual(selected.shape, (samples, features))
+        np.testing.assert_array_equal(selected, values[:, :, 2])
+
+    def test_3d_array_multiclass_without_class_raises_clear_error(self):
+        values = np.zeros((4, 2, 5))
+        with self.assertRaises(ValueError) as ctx:
+            select_shap_values(values)
+        self.assertIn("class_of_interest", str(ctx.exception))
+
+    def test_class_of_interest_out_of_bounds_raises(self):
+        values = np.zeros((4, 2, 3))
+        with self.assertRaises(ValueError):
+            select_shap_values(values, class_of_interest=9)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/validmind/tests/model_validation/sklearn/CalibrationCurve.py b/validmind/tests/model_validation/sklearn/CalibrationCurve.py
index fea0211b4..b3075c2fc 100644
--- a/validmind/tests/model_validation/sklearn/CalibrationCurve.py
+++ b/validmind/tests/model_validation/sklearn/CalibrationCurve.py
@@ -4,10 +4,12 @@
 
 from typing import Tuple
 
+import numpy as np
 import plotly.graph_objects as go
 from sklearn.calibration import calibration_curve
 
 from validmind import tags, tasks
+from validmind.errors import SkipTestError
 from validmind.vm_models import VMDataset, VMModel
 from validmind.vm_models.result import RawData
 
@@ -70,6 +72,14 @@ def CalibrationCurve(
     - Assumes bin boundaries are appropriate for the problem
     - May be affected by class imbalance
     """
+    # Binary-only by design: sklearn's calibration_curve raises a cryptic
+    # "pos_label is not specified" error on multiclass targets, so skip cleanly
+    # like ROCCurve/PrecisionRecallCurve rather than crashing.
+    if len(np.unique(dataset.y)) > 2:
+        raise SkipTestError(
+            "Calibration Curve is only supported for binary classification models"
+        )
+
     prob_true, prob_pred = calibration_curve(
         dataset.y, dataset.y_prob(model), n_bins=n_bins
     )
diff --git a/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py b/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py
index 6d6ba4936..c369027c5 100644
--- a/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py
+++ b/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py
@@ -31,6 +31,23 @@
 logger = get_logger(__name__)
 
 
+def _resolve_class_index(num_classes: int, class_of_interest: Optional[int]) -> int:
+    """Pick which class's SHAP values to use.
+
+    Defaults to class 1 for binary classification (matching the historical
+    behavior) and otherwise requires an explicit, in-bounds ``class_of_interest``.
+    """
+    if num_classes == 2 and class_of_interest is None:
+        return 1
+    if class_of_interest is not None and 0 <= class_of_interest < num_classes:
+        return class_of_interest
+    raise ValueError(
+        f"Invalid class_of_interest: {class_of_interest}. Must be between 0 and "
+        f"{num_classes - 1}. Multiclass SHAP importance requires an explicit "
+        "class_of_interest."
+    )
+
+
 def select_shap_values(
     shap_values: Union[np.ndarray, List[np.ndarray]],
     class_of_interest: Optional[int] = None,
@@ -40,9 +57,11 @@ def select_shap_values(
     For regression models, returns the SHAP values directly as there are no classes.
 
     Args:
-        shap_values: The SHAP values returned by the SHAP explainer. For multiclass
-            classification, this will be a list where each element corresponds to a class.
-            For regression, this will be a single array of SHAP values.
+        shap_values: The SHAP values returned by the SHAP explainer. Depending on the
+            SHAP version, multiclass classification is returned either as a list with
+            one ``(samples, features)`` array per class or as a single
+            ``(samples, features, classes)`` array. Binary classification and
+            regression are returned as a single ``(samples, features)`` array.
         class_of_interest: The class index for which to retrieve SHAP values. If None
             (default), the function will assume binary classification and use class 1
             by default.
@@ -52,24 +71,26 @@ def select_shap_values(
         output.
 
     Raises:
-        ValueError: If class_of_interest is specified and is out of bounds for the
-            number of classes.
+        ValueError: If class_of_interest is out of bounds, or is not specified for a
+            multiclass model.
     """
-    if not isinstance(shap_values, list):
-        # For regression, return the SHAP values as they are
-        selected_values = shap_values
-    else:
+    if isinstance(shap_values, list):
+        # Older SHAP returns a list with one (samples, features) array per class.
         num_classes = len(shap_values)
-        # Default to class 1 for binary classification where no class is specified
-        if num_classes == 2 and class_of_interest is None:
-            selected_values = shap_values[1]
-        # Otherwise, use the specified class_of_interest
-        elif class_of_interest is not None and 0 <= class_of_interest < num_classes:
-            selected_values = shap_values[class_of_interest]
-        else:
-            raise ValueError(
-                f"Invalid class_of_interest: {class_of_interest}. Must be between 0 and {num_classes - 1}."
-            )
+        selected_values = shap_values[
+            _resolve_class_index(num_classes, class_of_interest)
+        ]
+    elif isinstance(shap_values, np.ndarray) and shap_values.ndim == 3:
+        # Newer SHAP returns a single (samples, features, classes) array for
+        # multiclass; select along the trailing class axis so downstream plotting
+        # receives a 2D (samples, features) array instead of raising an IndexError.
+        num_classes = shap_values.shape[-1]
+        selected_values = shap_values[
+            :, :, _resolve_class_index(num_classes, class_of_interest)
+        ]
+    else:
+        # Regression, or binary returned as a single (samples, features) array.
+        selected_values = shap_values
 
     # Add type conversion here to ensure proper float array
     if hasattr(selected_values, "dtype"):