From d678e0911afc6b7654842fc5450e2fb5997356f7 Mon Sep 17 00:00:00 2001
From: hassaanch23 <imhassaan04@gmail.com>
Date: Sat, 29 Aug 2026 13:51:00 +0500
Subject: [PATCH] fix(r_squared): implement the documented zero_division
 parameter

`_KWARGS_DESCRIPTION` documents a `zero_division` argument:

    zero_division: Which value to substitute as a metric value when
        encountering zero division. Should be one of 0, 1, "warn". "warn"
        acts as 0, but the warning is raised.

but `_compute(self, predictions=None, references=None)` never accepted it, so
passing it raised TypeError, and the case it exists to control was unhandled.

When every reference is identical there is no variance to explain, so the sum
of squared total is zero and R^2 is undefined. Left to numpy the division
returned -inf, or nan when the predictions matched exactly, with only a
RuntimeWarning. Both propagate silently through any downstream aggregation:
one constant-reference batch turns an averaged score into nan.

    predictions=[1,2,3,4], references=[5,5,5,5]   before -inf   sklearn 0.0
    predictions=[5,5,5,5], references=[5,5,5,5]   before  nan   sklearn 1.0

`zero_division` now works as documented, defaulting to "warn" so the
previously silent case becomes visible rather than changing quietly to a
number the caller did not choose. An invalid value raises ValueError instead
of being ignored.

Non-degenerate inputs are untouched and still agree with
sklearn.metrics.r2_score.

Adds doctest examples for both settings, which is how metrics in this repo are
covered; tests/test_metric_common.py passes for r_squared.
---
 metrics/r_squared/r_squared.py | 36 +++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/metrics/r_squared/r_squared.py b/metrics/r_squared/r_squared.py
index 75e0f7e00..62dc28e2e 100644
--- a/metrics/r_squared/r_squared.py
+++ b/metrics/r_squared/r_squared.py
@@ -15,6 +15,8 @@
 """R squared metric."""
 
 
+import warnings
+
 import datasets
 import numpy as np
 
@@ -61,6 +63,14 @@
     >>> r_squared = r2_metric.compute(predictions=[1, 2, 3, 4], references=[0.9, 2.1, 3.2, 3.8])
     >>> print(r_squared)
     0.98
+
+    R^2 is undefined when the references have zero variance, since there is no
+    variance to explain. `zero_division` selects what to return instead:
+
+    >>> r2_metric.compute(predictions=[1, 2, 3, 4], references=[5, 5, 5, 5], zero_division=0)
+    0.0
+    >>> r2_metric.compute(predictions=[5, 5, 5, 5], references=[5, 5, 5, 5], zero_division=1)
+    1.0
 """
 
 
@@ -83,17 +93,25 @@ def _info(self):
             ],
         )
 
-    def _compute(self, predictions=None, references=None):
+    def _compute(self, predictions=None, references=None, zero_division="warn"):
         """
         Computes the coefficient of determination (R-squared) of predictions with respect to references.
 
         Parameters:
             predictions (List or np.ndarray): The predicted values.
             references (List or np.ndarray): The true/reference values.
+            zero_division (0, 1 or "warn"): Value to return when `references` has zero
+                variance, which makes the sum of squared total zero and R^2 undefined.
+                "warn" returns 0.0 and raises a warning.
 
         Returns:
             float: The R-squared value, rounded to 3 decimal places.
         """
+        if zero_division not in (0, 1, "warn"):
+            raise ValueError(
+                f'zero_division must be one of 0, 1 or "warn", got {zero_division!r}'
+            )
+
         predictions = np.array(predictions)
         references = np.array(references)
 
@@ -106,6 +124,22 @@ def _compute(self, predictions=None, references=None):
         # Calculate sum of squared total
         sst = np.sum((references - mean_references) ** 2)
 
+        # R^2 is undefined when the references never vary: there is no variance to
+        # explain, so the ratio divides by zero. Left to numpy this returns -inf
+        # (or nan when the predictions match exactly), which then propagates
+        # silently through any downstream aggregation.
+        if sst == 0:
+            if zero_division == "warn":
+                warnings.warn(
+                    "R^2 is undefined when all `references` are identical (zero "
+                    "variance); returning 0.0. Pass zero_division=0 or 1 to choose "
+                    "the value and silence this warning.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                return 0.0
+            return float(zero_division)
+
         # Calculate R Squared
         r_squared = 1 - (ssr / sst)