Merge remote-tracking branch 'origin/master' into fix/torch-export-preserve-dims

Author: MilesCranmerBot

.github/workflows/update_backend.yml

@@ -20,9 +20,7 @@ jobs:
           cache: pip
 
       - name: "Install dependencies"
-        run: |
-          python -m pip install --upgrade pip
-          pip install tomlkit
+        run: python -m pip install --upgrade pip
 
       - name: "Get SymbolicRegression.jl latest version"
         id: get-latest
@@ -36,13 +34,6 @@ jobs:
         run: |
           python .github/workflows/update_backend_version.py ${{ steps.get-latest.outputs.version }}
 
-      - name: "Restore changes if no diff to `pysr/juliapkg.json`"
-        run: |
-          if git diff --quiet pysr/juliapkg.json; then
-            echo "No changes to pysr/juliapkg.json. Restoring changes."
-            git restore pyproject.toml .release-please-manifest.json
-          fi
-
       - name: "Create PR if necessary"
         id: cpr
         uses: peter-evans/create-pull-request@v8
@@ -56,8 +47,6 @@ jobs:
           delete-branch: true
           commit-message: "chore: update backend to v${{ steps.get-latest.outputs.version }}"
           add-paths: |
-            .release-please-manifest.json
-            pyproject.toml
             pysr/juliapkg.json
 
       - name: "Trigger CI workflows (backend update PR)"

.github/workflows/update_backend_version.py

@@ -1,65 +1,17 @@
 import json
-import re
 import sys
 from pathlib import Path
 
-import tomlkit
-
 new_backend_version = sys.argv[1]
 
 assert not new_backend_version.startswith("v"), "Version should not start with 'v'"
 
 repo_root = Path(__file__).parent / ".." / ".."
-pyproject_toml = repo_root / "pyproject.toml"
 juliapkg_json = repo_root / "pysr" / "juliapkg.json"
-release_please_manifest = repo_root / ".release-please-manifest.json"
-
-with open(pyproject_toml) as toml_file:
-    pyproject_data = tomlkit.parse(toml_file.read())
-
 with open(juliapkg_json) as f:
     juliapkg_data = json.load(f)
 
-with open(release_please_manifest) as f:
-    release_please_manifest_data = json.load(f)
-
-current_version = pyproject_data["project"]["version"]
-parts = current_version.split(".")
-
-if len(parts) < 3:
-    raise ValueError(
-        f"Invalid version format: {current_version}. Expected at least 3 components (major.minor.patch)"
-    )
-
-major, minor = parts[0], parts[1]
-
-patch_match = re.match(r"^(\d+)(.*)$", parts[2])
-if not patch_match:
-    raise ValueError(
-        f"Could not parse patch version from '{parts[2]}' in version {current_version}. "
-        f"Expected patch to start with a number (e.g., '0', '1a1', '2rc3')"
-    )
-
-patch_num_str, patch_suffix = patch_match.groups()
-patch_num = int(patch_num_str)
-
-pre_release_match = re.fullmatch(r"(a|b|rc)(\d+)", patch_suffix)
-if pre_release_match:
-    pre_tag, pre_num = pre_release_match.groups()
-    new_patch = patch_num
-    new_suffix = f"{pre_tag}{int(pre_num) + 1}"
-else:
-    new_patch = patch_num + 1
-    new_suffix = patch_suffix
-
-# Add back any additional version components (e.g., "2.0.0.dev1" -> ".dev1")
-extra_parts = "." + ".".join(parts[3:]) if len(parts) > 3 else ""
-new_version = f"{major}.{minor}.{new_patch}{new_suffix}{extra_parts}"
-
-pyproject_data["project"]["version"] = new_version
-release_please_manifest_data["."] = new_version
-
-# Update backend - maintain current format (either "rev" or "version")
+# Update backend, maintain current format (either "rev" or "version")
 backend_pkg = juliapkg_data["packages"]["SymbolicRegression"]
 if "rev" in backend_pkg:
     backend_pkg["rev"] = f"v{new_backend_version}"
@@ -70,13 +22,6 @@
         "SymbolicRegression package must have either 'rev' or 'version' field"
     )
 
-with open(pyproject_toml, "w") as toml_file:
-    toml_file.write(tomlkit.dumps(pyproject_data))
-
 with open(juliapkg_json, "w") as f:
     json.dump(juliapkg_data, f, indent=4)
     f.write("\n")
-
-with open(release_please_manifest, "w") as f:
-    json.dump(release_please_manifest_data, f, indent=2)
-    f.write("\n")

.release-please-manifest.json

@@ -1,3 +1,3 @@
 {
-  ".": "2.0.0a2"
+  ".": "2.0.0a1"
 }

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)