fix: correct type in elementwise loss validation (MilesCranmer#1184)

* fix: warn on elementwise loss probe mismatch

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: simplify elementwise loss probe dtype

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

* Revert "fix: simplify elementwise loss probe dtype"

This reverts commit 775ab95.

* fix: keep elementwise loss mismatch strict

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

* docs: clarify elementwise loss validator docstring

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

* refactor: rename loss probe parameter

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

* test: update loss probe helper call

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

---------

Co-authored-by: MilesCranmerBot <milescranmerbot@users.noreply.github.com>
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

pysr/sr.py

@@ -235,26 +235,34 @@ def _check_assertions(
             )
 
 
-def _validate_elementwise_loss(custom_loss, *, has_weights: bool) -> None:
-    """Validate that a Julia `elementwise_loss` is callable.
+def _validate_elementwise_loss(
+    custom_loss, *, has_weights: bool, probe_value: Any = 1.0
+) -> None:
+    """Check whether a Julia `elementwise_loss` accepts the expected inputs.
 
-    We require exactly 2 args unless the user passed `weights=` to fit,
-    in which case we require 3 args.
+    The function probes the loss with two or three arguments, depending on
+    whether weights are present, using the same dtype that fitting will use.
+    If the probe fails, it raises a `ValueError` describing the expected
+    signature.
     """
 
     # 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 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`."
             )
     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 "
@@ -2109,13 +2117,19 @@ 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))
+
         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=np_dtype(1.0),
+            )
 
         custom_full_objective = jl.seval(
             str(self.loss_function) if self.loss_function is not None else "nothing"
@@ -2304,8 +2318,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

@@ -313,6 +313,34 @@ def test_elementwise_loss_with_weights_accepts_three_args(self):
         weights = np.array([1.0, 1.0])
         model.fit(X, y, weights=weights)
 
+    def test_elementwise_loss_float32_probe_accepts_strictly_typed_loss(self):
+        custom_loss = jl.seval(
+            "(prediction::Float32, target::Float32) -> (prediction - target)^2"
+        )
+        _validate_elementwise_loss(
+            custom_loss,
+            has_weights=False,
+            probe_value=np.float32(1.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=(
+                "(prediction::Float32, target::Float32) -> (prediction - target)^2"
+            ),
+        )
+        X = np.array([[0.0], [1.0]], dtype=np.float32)
+        y = np.array([0.0, 1.0], dtype=np.float32)
+        model.fit(X, y)
+
     def test_validation_helpers_skip_nonfunction(self):
         _validate_elementwise_loss(jl.seval("1.0"), has_weights=False)