fix: clarify loss validator arity + error text

Author: MilesCranmerBot

pysr/sr.py

@@ -251,7 +251,7 @@ def _validate_elementwise_loss(custom_loss, *, has_weights: bool) -> None:
         ok = bool(jl.applicable(custom_loss, 1.0, 1.0, 1.0))
         if not ok:
             raise ValueError(
-                "`elementwise_loss` must accept (prediction, target, weights) when `weights` is passed to `fit`."
+                "`elementwise_loss` must accept (prediction, target, weight) when `weights` is passed to `fit`."
             )
     else:
         ok = bool(jl.applicable(custom_loss, 1.0, 1.0))
@@ -274,16 +274,17 @@ def _validate_custom_objective(
         raise ValueError(f"`{knob}` must evaluate to a callable Julia function.")
 
     methods = jl.collect(jl.methods(custom_objective))
-    accepts_three_args = any(
-        (not bool(m.isva) and int(m.nargs) == 4) or (bool(m.isva) and int(m.nargs) <= 4)
-        for m in methods
-    )
-    appears_elementwise = any(
-        (not bool(m.isva) and int(m.nargs) == 3) or (bool(m.isva) and int(m.nargs) <= 3)
-        for m in methods
-    )
 
-    if not accepts_three_args and appears_elementwise:
+    def _accepts_npos(m, npos: int) -> bool:
+        required_npos = int(m.nargs) - 1
+        if bool(m.isva):
+            return required_npos <= npos
+        return required_npos == npos
+
+    accepts_three_args = any(_accepts_npos(m, 3) for m in methods)
+    accepts_two_args = any(_accepts_npos(m, 2) for m in methods)
+
+    if not accepts_three_args and accepts_two_args:
         msg = (
             f"`{knob}` must have signature like {signature}. "
             f"If you intended an elementwise loss, use `{elementwise_alternative}`."

pysr/test/test_main.py

@@ -242,6 +242,26 @@ def test_loss_function_valid_full_objective_runs(self):
         y = np.array([0.0, 1.0])
         model.fit(X, y)
 
+    def test_loss_function_varargs_objective_runs(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function="""
+            begin
+                varloss(tree, dataset, options...) = zero(eltype(dataset.y))
+                varloss
+            end
+            """,
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        model.fit(X, y)
+
     def test_elementwise_loss_wrong_signature_errors_early(self):
         """Validate `elementwise_loss` signature (prediction, target[, weights])."""
         model = PySRRegressor(