Only install packages as required

pysr/julia_extensions.py

@@ -0,0 +1,25 @@
+"""This file installs and loads extensions for SymbolicRegression."""
+from .julia_import import jl
+
+
+def load_required_packages(*, turbo=False, enable_autodiff=False):
+    if turbo:
+        load_package("LoopVectorization")
+    if enable_autodiff:
+        load_package("Zygote")
+    if cluster_manager is not None:
+        load_package("ClusterManagers")
+
+
+def load_package(package_name):
+    jl.seval(f"""
+        try
+            using {package_name}
+        catch e
+            isa(e, ArgumentError) || throw(e)
+            using Pkg: Pkg
+            Pkg.add("{package_name}")
+            using {package_name}
+        end
+    """)
+    return None

pysr/juliapkg.json

@@ -5,17 +5,9 @@
             "uuid": "8254be44-1295-4e6a-a16d-46603ac705cb",
             "version": "=0.24.0"
         },
-        "ClusterManagers": {
-            "uuid": "34f1f09b-3a8b-5176-ab39-66d58a4d544e",
-            "version": "0.4"
-        },
         "Serialization": {
             "uuid": "9e88b42a-f829-5b0c-bbe9-9e923198166b",
             "version": "1"
-        },
-        "Zygote": {
-            "uuid": "e88e6eb3-aa80-5325-afca-941959d7151f",
-            "version": "0.6"
         }
     }
 }

pysr/sr.py

@@ -32,6 +32,7 @@ from .export_numpy import sympy2numpy
 from .export_sympy import assert_valid_sympy_symbol, create_sympy_symbols, pysr2sympy
 from .export_torch import sympy2torch
 from .feature_selection import run_feature_selection
+from .julia_extensions import load_required_packages
 from .julia_helpers import (
     PythonCall,
     _escape_filename,
@@ -1605,6 +1606,12 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             else "nothing"
         )
 
+        load_required_packages(
+            turbo=turbo,
+            enable_autodiff=enable_autodiff,
+            cluster_manager=cluster_manager
+        )
+
         mutation_weights = SymbolicRegression.MutationWeights(
             mutate_constant=self.weight_mutate_constant,
             mutate_operator=self.weight_mutate_operator,