fix: simplify elementwise loss probe dtype

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

Author: MilesCranmerBot

pysr/sr.py

@@ -236,7 +236,7 @@ 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.
 
@@ -251,6 +251,7 @@ def _validate_elementwise_loss(
     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
@@ -2125,7 +2126,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 +2136,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

@@ -265,7 +265,7 @@ def test_elementwise_loss_wrong_signature_warns(self):
             _validate_elementwise_loss(
                 custom_loss,
                 has_weights=False,
-                probe_value=np.float32(1.0),
+                probe_dtype=np.float32,
             )
 
     def test_elementwise_loss_with_weights_requires_three_args_warns(self):
@@ -274,7 +274,7 @@ def test_elementwise_loss_with_weights_requires_three_args_warns(self):
             _validate_elementwise_loss(
                 custom_loss,
                 has_weights=True,
-                probe_value=np.float32(1.0),
+                probe_dtype=np.float32,
             )
 
     def test_elementwise_loss_with_weights_accepts_three_args(self):
@@ -304,7 +304,7 @@ def test_elementwise_loss_float32_probe_accepts_strictly_typed_loss(self):
             _validate_elementwise_loss(
                 custom_loss,
                 has_weights=False,
-                probe_value=np.float32(1.0),
+                probe_dtype=np.float32,
             )
         self.assertEqual(len(caught), 0)