fix: warn on elementwise loss probe mismatch

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

Author: MilesCranmerBot

pysr/sr.py

@@ -235,31 +235,46 @@ def _check_assertions(
             )
 
 
-def _validate_elementwise_loss(custom_loss, *, has_weights: bool) -> None:
+def _validate_elementwise_loss(
+    custom_loss, *, has_weights: bool, probe_value: Any = 1.0
+) -> None:
     """Validate that a Julia `elementwise_loss` is callable.
 
-    We require exactly 2 args unless the user passed `weights=` to fit,
-    in which case we require 3 args.
+    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_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:
-        ok = bool(jl.applicable(custom_loss, 1.0, 1.0, 1.0))
-        if not ok:
-            raise ValueError(
-                "`elementwise_loss` must accept (prediction, target, weight) when `weights` is passed to `fit`."
-            )
+        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,
+        )
     else:
-        ok = bool(jl.applicable(custom_loss, 1.0, 1.0))
-        if not ok:
-            raise ValueError(
-                "`elementwise_loss` must accept (prediction, target). If you intended a full objective, use "
-                "`loss_function` or `loss_function_expression`."
-            )
+        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,
+        )
 
 
 def _validate_custom_objective(
@@ -2109,13 +2124,20 @@ def _run(
         if isinstance(complexity_of_variables, list):
             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)
             if self.elementwise_loss is not None
             else "nothing"
         )
         if self.elementwise_loss is not None:
-            _validate_elementwise_loss(custom_loss, has_weights=weights is not None)
+            _validate_elementwise_loss(
+                custom_loss,
+                has_weights=weights is not None,
+                probe_value=probe_value,
+            )
 
         custom_full_objective = jl.seval(
             str(self.loss_function) if self.loss_function is not None else "nothing"
@@ -2304,8 +2326,6 @@ def _run(
         self.julia_options_stream_ = jl_serialize(options)
 
         # Convert data to desired precision
-        test_X = np.array(X)
-        np_dtype = self._get_precision_mapped_dtype(test_X)
 
         # This converts the data into a Julia array:
         jl_X = jl_array(np.array(X, dtype=np_dtype).T)

pysr/test/test_main.py

@@ -258,25 +258,28 @@ 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_errors_early(self):
+    def test_elementwise_loss_wrong_signature_warns(self):
         """Validate `elementwise_loss` signature (prediction, target[, weights])."""
-        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",
+        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),
+            )
+
+    def test_elementwise_loss_with_weights_requires_three_args_warns(self):
+        custom_loss = jl.seval(
+            "myloss2(prediction, target) = (prediction - target)^2"
         )
-        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))
+        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):
+    def test_elementwise_loss_with_weights_accepts_three_args(self):
         model = PySRRegressor(
             niterations=1,
             populations=1,
@@ -285,33 +288,51 @@ def test_elementwise_loss_with_weights_requires_three_args(self):
             verbosity=0,
             temp_equation_file=True,
             binary_operators=["+"],
-            elementwise_loss="myloss2(prediction, target) = (prediction - target)^2",
+            elementwise_loss=(
+                "myloss3(prediction, target, weights) = weights * (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))
+        model.fit(X, y, weights=weights)
 
-    def test_elementwise_loss_with_weights_accepts_three_args(self):
+    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)
+
+    def test_elementwise_loss_float32_fit_accepts_strictly_typed_loss(self):
         model = PySRRegressor(
             niterations=1,
             populations=1,
             procs=0,
             progress=False,
             verbosity=0,
+            precision=32,
             temp_equation_file=True,
             binary_operators=["+"],
             elementwise_loss=(
-                "myloss3(prediction, target, weights) = weights * (prediction - target)^2"
+                "(prediction::Float32, target::Float32) -> (prediction - target)^2"
             ),
         )
-        X = np.array([[0.0], [1.0]])
-        y = np.array([0.0, 1.0])
-        weights = np.array([1.0, 1.0])
-        model.fit(X, y, weights=weights)
+        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],
+        )
 
     def test_validation_helpers_skip_nonfunction(self):
         _validate_elementwise_loss(jl.seval("1.0"), has_weights=False)