Ensure julia is loaded as soon as possible

pysr/__init__.py

@@ -1,8 +1,13 @@
+# 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 . import sklearn_monkeypatch
 from .deprecated import best, best_callable, best_row, best_tex, pysr
 from .export_jax import sympy2jax
 from .export_torch import sympy2torch
-from .sr import PySRRegressor, jl
+from .sr import PySRRegressor
 
 # This file is created by setuptools_scm during the build process:
 from .version import __version__
@@ -20,4 +25,5 @@
     "best_tex",
     "pysr",
     "__version__",
+    "jl",
 ]

pysr/julia_helpers.py

@@ -1,39 +1,8 @@
 """Functions for initializing the Julia environment and installing deps."""
-import os
-import sys
-import warnings
-
-if "juliacall" in sys.modules:
-    warnings.warn(
-        "juliacall module already imported. Make sure that you have set `PYTHON_JULIACALL_HANDLE_SIGNALS=yes` to avoid segfaults."
-    )
-
-# Required to avoid segfaults (https://juliapy.github.io/PythonCall.jl/dev/faq/)
-if os.environ.get("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes") != "yes":
-    warnings.warn(
-        "PYTHON_JULIACALL_HANDLE_SIGNALS environment variable is set to something other than 'yes' or ''. "
-        + "You will experience segfaults if running with multithreading."
-    )
-
-if os.environ.get("JULIA_NUM_THREADS", "auto") != "auto":
-    warnings.warn(
-        "JULIA_NUM_THREADS environment variable is set to something other than 'auto', "
-        "so PySR was not able to set it. You may wish to set it to `'auto'` for full use "
-        "of your CPU."
-    )
-
-# TODO: Remove these when juliapkg lets you specify this
-for k, default in (
-    ("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes"),
-    ("JULIA_NUM_THREADS", "auto"),
-    ("JULIA_OPTIMIZE", "3"),
-):
-    os.environ[k] = os.environ.get(k, default)
-
-
-from juliacall import Main as jl  # type: ignore
 from juliacall import convert as jl_convert  # type: ignore
 
+from .julia_import import jl
+
 jl.seval("using Serialization: Serialization")
 jl.seval("using PythonCall: PythonCall")
 

pysr/julia_import.py

@@ -0,0 +1,32 @@
+import os
+import sys
+import warnings
+
+if "juliacall" in sys.modules:
+    warnings.warn(
+        "juliacall module already imported. Make sure that you have set `PYTHON_JULIACALL_HANDLE_SIGNALS=yes` to avoid segfaults."
+    )
+
+# Required to avoid segfaults (https://juliapy.github.io/PythonCall.jl/dev/faq/)
+if os.environ.get("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes") != "yes":
+    warnings.warn(
+        "PYTHON_JULIACALL_HANDLE_SIGNALS environment variable is set to something other than 'yes' or ''. "
+        + "You will experience segfaults if running with multithreading."
+    )
+
+if os.environ.get("JULIA_NUM_THREADS", "auto") != "auto":
+    warnings.warn(
+        "JULIA_NUM_THREADS environment variable is set to something other than 'auto', "
+        "so PySR was not able to set it. You may wish to set it to `'auto'` for full use "
+        "of your CPU."
+    )
+
+# TODO: Remove these when juliapkg lets you specify this
+for k, default in (
+    ("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes"),
+    ("JULIA_NUM_THREADS", "auto"),
+    ("JULIA_OPTIMIZE", "3"),
+):
+    os.environ[k] = os.environ.get(k, default)
+
+from juliacall import Main as jl  # type: ignore

pysr/sr.py

@@ -35,10 +35,10 @@
 from .julia_helpers import (
     _escape_filename,
     _load_cluster_manager,
-    jl,
     jl_array,
     jl_deserialize_s,
 )
+from .julia_import import jl
 from .utils import (
     _csv_filename_to_pkl_filename,
     _preprocess_julia_floats,