Move JAX export to separate file

pysr/__init__.py

@@ -1,3 +1,3 @@
 from .sr import pysr, get_hof, best, best_tex, best_callable, best_row
 from .feynman_problems import Problem, FeynmanProblem
-from .export import sympy2jax
+from .export_jax import sympy2jax

pysr/{export.py → export_jax.py}

@@ -75,7 +75,7 @@ def sympy2jaxtext(expr, parameters, symbols_in):
         else:
             return f'{_func}({", ".join(args)})'
 
-def sympy2jax(equation, symbols_in):
+def sympy2jax(expression, symbols_in):
     """Returns a function f and its parameters;
     the function takes an input matrix, and a list of arguments:
             f(X, parameters)
@@ -147,8 +147,8 @@ def sympy2jax(equation, symbols_in):
         ```
     """
     parameters = []
-    functional_form_text = sympy2jaxtext(equation, parameters, symbols_in)
-    hash_string = 'A_' + str(abs(hash(str(equation) + str(symbols_in))))
+    functional_form_text = sympy2jaxtext(expression, parameters, symbols_in)
+    hash_string = 'A_' + str(abs(hash(str(expression) + str(symbols_in))))
     text = f"def {hash_string}(X, parameters):\n"
     text += "    return "
     text += functional_form_text

pysr/sr.py

@@ -13,7 +13,7 @@ import shutil
 from pathlib import Path
 from datetime import datetime
 import warnings
-from .export import sympy2jax
+from .export_jax import sympy2jax
 
 global_equation_file = 'hall_of_fame.csv'
 global_n_features = None