Export SymbolicRegression.jl module

pysr/__init__.py

@@ -1,7 +1,7 @@
 # This must be imported as early as possible to prevent
 # library linking issues caused by numpy/pytorch/etc. importing
 # old libraries:
-from .julia_import import jl  # isort:skip
+from .julia_import import jl, SymbolicRegression  # isort:skip
 
 from . import sklearn_monkeypatch
 from .deprecated import best, best_callable, best_row, best_tex, pysr
@@ -14,6 +14,8 @@ from .sr import PySRRegressor
 from .version import __version__
 
 __all__ = [
+    "jl",
+    "SymbolicRegression",
     "sklearn_monkeypatch",
     "sympy2jax",
     "sympy2torch",
@@ -26,6 +28,5 @@ __all__ = [
     "best_tex",
     "pysr",
     "__version__",
-    "jl",
     "install",
 ]

pysr/julia_import.py

@@ -66,3 +66,6 @@ if jl.seval('VERSION >= v"1.9.0-DEV.0"'):
         return jl.eval(jl.Meta.parseall(s))
 
     jl.seval = seval
+
+jl.seval("using SymbolicRegression")
+SymbolicRegression = jl.SymbolicRegression

pysr/sr.py

@@ -38,7 +38,7 @@ from .julia_helpers import (
     jl_array,
     jl_deserialize_s,
 )
-from .julia_import import jl
+from .julia_import import SymbolicRegression, jl
 from .utils import (
     _csv_filename_to_pkl_filename,
     _preprocess_julia_floats,
@@ -1550,9 +1550,6 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         if cluster_manager is not None:
             cluster_manager = _load_cluster_manager(cluster_manager)
 
-        jl.seval("using SymbolicRegression")
-        SymbolicRegression = jl.SymbolicRegression
-
         jl.plus = jl.seval("(+)")
         jl.sub = jl.seval("(-)")
         jl.mult = jl.seval("(*)")