extreme_stable case for mean_tweedie_deviance (scikit-learn#29258)

Author: EmilyXinyi

sklearn/metrics/_regression.py

@@ -1284,12 +1284,14 @@ def _mean_tweedie_deviance(y_true, y_pred, sample_weight, power):
     """Mean Tweedie deviance regression loss."""
     xp, _ = get_namespace(y_true, y_pred)
     p = power
+    zero = xp.asarray(0, dtype=y_true.dtype)
     if p < 0:
         # 'Extreme stable', y any real number, y_pred > 0
         dev = 2 * (
-            xp.pow(xp.where(y_true > 0, y_true, 0), 2 - p) / ((1 - p) * (2 - p))
-            - y_true * xp.pow(y_pred, 1 - p) / (1 - p)
-            + xp.pow(y_pred, 2 - p) / (2 - p)
+            xp.pow(xp.where(y_true > 0, y_true, zero), xp.asarray(2 - p))
+            / ((1 - p) * (2 - p))
+            - y_true * xp.pow(y_pred, xp.asarray(1 - p)) / (1 - p)
+            + xp.pow(y_pred, xp.asarray(2 - p)) / (2 - p)
         )
     elif p == 0:
         # Normal distribution, y and y_pred any real number
@@ -1302,9 +1304,9 @@ def _mean_tweedie_deviance(y_true, y_pred, sample_weight, power):
         dev = 2 * (xp.log(y_pred / y_true) + y_true / y_pred - 1)
     else:
         dev = 2 * (
-            xp.pow(y_true, 2 - p) / ((1 - p) * (2 - p))
-            - y_true * xp.pow(y_pred, 1 - p) / (1 - p)
-            + xp.pow(y_pred, 2 - p) / (2 - p)
+            xp.pow(y_true, xp.asarray(2 - p)) / ((1 - p) * (2 - p))
+            - y_true * xp.pow(y_pred, xp.asarray(1 - p)) / (1 - p)
+            + xp.pow(y_pred, xp.asarray(2 - p)) / (2 - p)
         )
     return float(_average(dev, weights=sample_weight))
 
@@ -1384,14 +1386,14 @@ def mean_tweedie_deviance(y_true, y_pred, *, sample_weight=None, power=0):
     message = f"Mean Tweedie deviance error with power={power} can only be used on "
     if power < 0:
         # 'Extreme stable', y any real number, y_pred > 0
-        if (y_pred <= 0).any():
+        if xp.any(y_pred <= 0):
             raise ValueError(message + "strictly positive y_pred.")
     elif power == 0:
         # Normal, y and y_pred can be any real number
         pass
     elif 1 <= power < 2:
         # Poisson and compound Poisson distribution, y >= 0, y_pred > 0
-        if (y_true < 0).any() or (y_pred <= 0).any():
+        if xp.any(y_true < 0) or xp.any(y_pred <= 0):
             raise ValueError(message + "non-negative y and strictly positive y_pred.")
     elif power >= 2:
         # Gamma and Extreme stable distribution, y and y_pred > 0

sklearn/metrics/tests/test_common.py

@@ -1991,6 +1991,8 @@ def check_array_api_metric_pairwise(metric, array_namespace, device, dtype_name)
         check_array_api_multilabel_classification_metric,
     ],
     mean_tweedie_deviance: [check_array_api_regression_metric],
+    partial(mean_tweedie_deviance, power=-0.5): [check_array_api_regression_metric],
+    partial(mean_tweedie_deviance, power=1.5): [check_array_api_regression_metric],
     r2_score: [
         check_array_api_regression_metric,
         check_array_api_regression_metric_multioutput,