fix: keep elementwise loss mismatch strict

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: MilesCranmerBot

pysr/sr.py

@@ -236,45 +236,37 @@ def _check_assertions(
 
 
 def _validate_elementwise_loss(
-    custom_loss, *, has_weights: bool, probe_value: Any = 1.0
+    custom_loss, *, has_weights: bool, probe_dtype: Callable[[float], Any] = float
 ) -> None:
     """Validate that a Julia `elementwise_loss` is callable.
 
     We probe with the dtype that the Julia backend will use, which avoids
     falsely rejecting strictly typed losses such as `(::Float32, ::Float32)`.
-    If the probe still fails, emit a warning rather than raising so Julia can
-    surface the real `MethodError` during fitting for advanced custom losses.
     """
 
     # This can be either a LossFunctions.jl object (e.g. `L2DistLoss()`) or a Julia function.
     # Only validate arity when the evaluated object is actually a function.
     if not jl_is_function(custom_loss):
         return
 
+    probe_value = probe_dtype(1.0)
     probe_args = (
         (probe_value, probe_value, probe_value)
         if has_weights
         else (probe_value, probe_value)
     )
     ok = bool(jl.applicable(custom_loss, *probe_args))
-    if ok:
-        return
-
     if has_weights:
-        warnings.warn(
-            "`elementwise_loss` did not match the probed (prediction, target, weight) signature "
-            "for the dtype used during fitting. Continuing anyway so Julia can surface a more "
-            "specific `MethodError` if needed.",
-            stacklevel=2,
-        )
+        if not ok:
+            raise ValueError(
+                "`elementwise_loss` must accept (prediction, target, weight) when `weights` is passed to `fit`."
+            )
     else:
-        warnings.warn(
-            "`elementwise_loss` did not match the probed (prediction, target) signature for "
-            "the dtype used during fitting. If you intended a full objective, use "
-            "`loss_function` or `loss_function_expression`. Continuing anyway so Julia can "
-            "surface a more specific `MethodError` if needed.",
-            stacklevel=2,
-        )
+        if not ok:
+            raise ValueError(
+                "`elementwise_loss` must accept (prediction, target). If you intended a full objective, use "
+                "`loss_function` or `loss_function_expression`."
+            )
 
 
 def _validate_custom_objective(
@@ -2125,7 +2117,6 @@ def _run(
             complexity_of_variables = jl_array(complexity_of_variables)
 
         np_dtype = self._get_precision_mapped_dtype(np.array(X))
-        probe_value = np_dtype(1.0)
 
         custom_loss = jl.seval(
             str(self.elementwise_loss)
@@ -2136,7 +2127,7 @@ def _run(
             _validate_elementwise_loss(
                 custom_loss,
                 has_weights=weights is not None,
-                probe_value=probe_value,
+                probe_dtype=np_dtype,
             )
 
         custom_full_objective = jl.seval(

pysr/test/test_main.py

@@ -258,24 +258,42 @@ def test_loss_function_varargs_objective_runs(self):
         y = np.array([0.0, 1.0])
         model.fit(X, y)
 
-    def test_elementwise_loss_wrong_signature_warns(self):
+    def test_elementwise_loss_wrong_signature_errors_early(self):
         """Validate `elementwise_loss` signature (prediction, target[, weights])."""
-        custom_loss = jl.seval("myloss_bad_arity(a) = a")
-        with self.assertWarnsRegex(UserWarning, "elementwise_loss"):
-            _validate_elementwise_loss(
-                custom_loss,
-                has_weights=False,
-                probe_value=np.float32(1.0),
-            )
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            elementwise_loss="myloss_bad_arity(a) = a",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y)
+        self.assertIn("elementwise_loss", str(cm.exception))
 
-    def test_elementwise_loss_with_weights_requires_three_args_warns(self):
-        custom_loss = jl.seval("myloss2(prediction, target) = (prediction - target)^2")
-        with self.assertWarnsRegex(UserWarning, "elementwise_loss"):
-            _validate_elementwise_loss(
-                custom_loss,
-                has_weights=True,
-                probe_value=np.float32(1.0),
-            )
+    def test_elementwise_loss_with_weights_requires_three_args(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            elementwise_loss="myloss2(prediction, target) = (prediction - target)^2",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        weights = np.array([1.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y, weights=weights)
+        self.assertIn("elementwise_loss", str(cm.exception))
+        self.assertIn("weights", str(cm.exception))
 
     def test_elementwise_loss_with_weights_accepts_three_args(self):
         model = PySRRegressor(
@@ -299,14 +317,11 @@ def test_elementwise_loss_float32_probe_accepts_strictly_typed_loss(self):
         custom_loss = jl.seval(
             "(prediction::Float32, target::Float32) -> (prediction - target)^2"
         )
-        with warnings.catch_warnings(record=True) as caught:
-            warnings.simplefilter("always")
-            _validate_elementwise_loss(
-                custom_loss,
-                has_weights=False,
-                probe_value=np.float32(1.0),
-            )
-        self.assertEqual(len(caught), 0)
+        _validate_elementwise_loss(
+            custom_loss,
+            has_weights=False,
+            probe_dtype=np.float32,
+        )
 
     def test_elementwise_loss_float32_fit_accepts_strictly_typed_loss(self):
         model = PySRRegressor(
@@ -324,13 +339,7 @@ def test_elementwise_loss_float32_fit_accepts_strictly_typed_loss(self):
         )
         X = np.array([[0.0], [1.0]], dtype=np.float32)
         y = np.array([0.0, 1.0], dtype=np.float32)
-        with warnings.catch_warnings(record=True) as caught:
-            warnings.simplefilter("always")
-            model.fit(X, y)
-        self.assertFalse(
-            any("elementwise_loss" in str(w.message) for w in caught),
-            msg=[str(w.message) for w in caught],
-        )
+        model.fit(X, y)
 
     def test_validation_helpers_skip_nonfunction(self):
         _validate_elementwise_loss(jl.seval("1.0"), has_weights=False)