Skip to content

Add experimental Rust backend adapter#1242

Draft
gustavorubim wants to merge 1 commit into
astroautomata:masterfrom
gustavorubim:rust-backend-prototype
Draft

Add experimental Rust backend adapter#1242
gustavorubim wants to merge 1 commit into
astroautomata:masterfrom
gustavorubim:rust-backend-prototype

Conversation

@gustavorubim

@gustavorubim gustavorubim commented Jun 29, 2026

Copy link
Copy Markdown

Summary

Adds an experimental backend=rust path plus default backend=auto selection to PySR. backend=auto uses the optional Rust backend when the symbolic_regression_rs module from pysr-rust-backend is installed, otherwise it falls back to the full Julia backend.

This PR intentionally does not vendor the Rust workspace into PySR. It depends on the Python wrapper package from astroautomata/symbolic_regression.rs being released as pysr-rust-backend>=0.1.0,<0.2.0.

What changed

  • Add PySRRegressor(backend=auto | julia | rust) and route Rust searches through pysr.backends.rust.
  • Add lazy Julia imports so from pysr import PySRRegressor does not initialize juliacall.
  • Add pysr[rust] optional dependency metadata for pysr-rust-backend.
  • Add Rust backend adapter tests with a mocked wrapper plus a real-wrapper smoke test when symbolic_regression_rs is installed.
  • Add backend docs and a Julia-vs-Rust benchmark script.
  • Add CI coverage for the adapter tests without requiring the unreleased optional Rust package.

Follow-up to review comments

  • Default backend is now backend=auto, as suggested, so installing the Rust backend selects Rust unless backend=julia is set explicitly.
  • The companion wrapper PR was updated to rename the distribution to pysr-rust-backend, propagate seed, reject unknown options, validate operator arity, document unsupported features, and smoke-test publish wheels.
  • A separate draft vendored-connector PR will be opened for comparison, without replacing this split-package design.

Validation

  • python -m pysr test main,jax,torch,autodiff,cli,dev,startup,slurm,rust with the Rust wrapper uninstalled: 143 tests passed, 5 skipped.
  • python -m pysr test rust with local pysr-rust-backend installed: 14 tests passed.
  • python -m mypy --ignore-missing-imports pysr
  • python -m compileall -q pysr benchmarks/compare_backends.py
  • pyproject TOML parse check
  • GitHub Actions YAML parse check
  • benchmarks/compare_backends.py smoke run for Julia and Rust: Julia 18.547s, Rust 0.016s on the small one-repeat fixture.

Dependencies

Known blocker

pysr-rust-backend is not visible on PyPI yet. This should stay draft until the wrapper package is merged/released or until maintainers choose the in-repo connector design.

@gustavorubim

gustavorubim commented Jun 29, 2026

Copy link
Copy Markdown
Author

Fresh-env benchmark smoke

I ran a fresh virtualenv benchmark to verify the PySR Rust backend adapter against the Julia backend across 10 generated target equations.

Setup:

  • Fresh venv created under the system temp directory
  • Python: 3.12
  • Installed symbolic-regression-rs==0.1.0 by building the wrapper from a local checkout of the symbolic_regression.rs wrapper PR branch
  • Installed this PySR branch with pip install -e ".[rust]"
  • Confirmed import pysr did not load juliacall
  • Benchmark: 10 synthetic equations, 96 samples, 3 features, deterministic serial execution, same small search budget for both backends

Cold fresh-env pass:

Backend OK Total fit time Mean Median Mean MSE
Julia 10/10 96.061s 9.606s 0.265s 0.309
Rust 10/10 0.134s 0.013s 0.012s 0.554

Warm pass after Julia deps/precompile:

Backend OK Total fit time Mean Median Mean MSE
Julia 10/10 20.020s 2.002s 0.237s 0.309
Rust 10/10 0.136s 0.014s 0.012s 0.554

Notes:

  • The first Julia timing dominates the totals: 93.6s cold due Julia package install/precompile, and 17.9s warm due remaining Julia startup/JIT cost.
  • Excluding the first warm Julia run, Julia averaged about 0.237s per target; Rust averaged about 0.013s per target.
  • This is a smoke/timing benchmark for adapter behavior, not an accuracy claim. The Rust backend currently has a reduced feature/search surface and the wrapper package still needs to be released to PyPI before this PR can be marked ready.

@MilesCranmer

Copy link
Copy Markdown
Member

Small initial comments:

  • It would be useful to see a PR that does vendor the rust-python interface so I can see if the plumbing actually works or not. In fact I almost wonder if it's better to put the connector package in this repository rather than in the sr.rs repo.
  • From an interface perspective It might be useful if the default was backend="auto" and it chooses backend based on what the user installed. Then I won't accidentally install all of Julia if I forget that option in the PySRRegressor.

@gustavorubim

Copy link
Copy Markdown
Author

Updated this PR in response to the initial feedback:

  • backend now defaults to auto, selecting the Rust backend when pysr-rust-backend / symbolic_regression_rs is installed and falling back to Julia otherwise.
  • The optional dependency is now pysr-rust-backend to match the renamed wrapper distribution.
  • The companion wrapper PR now propagates seed, rejects unknown options, validates operator arity, documents unsupported features, and smoke-tests publish wheels.
  • I also opened a separate draft comparison PR with the connector package living inside this repo: Draft: add in-repo Rust backend connector package #1251

Validation rerun:

  • Full PySR suite with Rust wrapper uninstalled: 143 passed, 5 skipped.
  • Rust adapter tests with local pysr-rust-backend installed: 14 passed.
  • mypy, compileall, pyproject parse, CI YAML parse.
  • Small benchmarks/compare_backends.py smoke run for both Julia and Rust.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants