scikit-learn · OmarManzoor · Jun 7, 2024 · Jun 6, 2024 · Jun 6, 2024 · Jun 6, 2024
diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -102,6 +102,7 @@ Metrics
 -------

 - :func:`sklearn.metrics.accuracy_score`
+- :func:`sklearn.metrics.d2_tweedie_score`
 - :func:`sklearn.metrics.mean_absolute_error`
 - :func:`sklearn.metrics.mean_tweedie_deviance`
 - :func:`sklearn.metrics.pairwise.cosine_similarity`

diff --git a/doc/whats_new/v1.6.rst b/doc/whats_new/v1.6.rst
@@ -32,6 +32,7 @@ See :ref:`array_api` for more details.

 **Functions:**

+- :func:`sklearn.metrics.d2_tweedie_score` :pr:`29207` by :user:`Emily Chen <EmilyXinyi>`;
 - :func:`sklearn.metrics.mean_absolute_error` :pr:`27736` by :user:`Edoardo Abati <EdAbati>`;
 - :func:`sklearn.metrics.mean_squared_error` :pr:`29142` by :user:`Yaroslav Korobko <Tialo>`;
 - :func:`sklearn.metrics.mean_tweedie_deviance` :pr:`28106` by :user:`Thomas Li <lithomas1>`;

diff --git a/sklearn/metrics/_regression.py b/sklearn/metrics/_regression.py
@@ -1598,8 +1598,10 @@ def d2_tweedie_score(y_true, y_pred, *, sample_weight=None, power=0):
    >>> d2_tweedie_score(y_true, y_true, power=2)
    1.0
    """
+    xp, _ = get_namespace(y_true, y_pred)
+
    y_type, y_true, y_pred, _ = _check_reg_targets(
-        y_true, y_pred, None, dtype=[np.float64, np.float32]
+        y_true, y_pred, None, dtype=[xp.float64, xp.float32], xp=xp
    )
    if y_type == "continuous-multioutput":
        raise ValueError("Multioutput not supported in d2_tweedie_score")
@@ -1609,12 +1611,12 @@ def d2_tweedie_score(y_true, y_pred, *, sample_weight=None, power=0):
        warnings.warn(msg, UndefinedMetricWarning)
        return float("nan")

-    y_true, y_pred = np.squeeze(y_true), np.squeeze(y_pred)
+    y_true, y_pred = xp.squeeze(y_true, axis=1), xp.squeeze(y_pred, axis=1)
    numerator = mean_tweedie_deviance(
        y_true, y_pred, sample_weight=sample_weight, power=power
    )

-    y_avg = np.average(y_true, weights=sample_weight)
+    y_avg = _average(y_true, weights=sample_weight, xp=xp)
    denominator = _mean_tweedie_deviance(
        y_true, y_avg, sample_weight=sample_weight, power=power
    )

diff --git a/sklearn/metrics/tests/test_common.py b/sklearn/metrics/tests/test_common.py
@@ -1942,6 +1942,9 @@ def check_array_api_metric_pairwise(metric, array_namespace, device, dtype_name)
        check_array_api_regression_metric,
        check_array_api_multioutput_regression_metric,
    ],
+    d2_tweedie_score: [
+        check_array_api_regression_metric,
+    ],
 }