Test sympy export to jax

.github/workflows/CI.yml

@@ -58,6 +58,7 @@ jobs:
         run: |
             python -m pip install --upgrade pip
             pip install -r requirements.txt
+            pip install jax jaxlib # (optional import)
             python setup.py install
         shell: bash
       - name: "Run tests"

test/test.py

@@ -1,5 +1,9 @@
 import numpy as np
-from pysr import pysr
+from pysr import pysr, sympy2jax
+from jax import numpy as jnp
+from jax import random
+from jax import grad
+import sympy
 X = np.random.randn(100, 5)
 
 print("Test 1 - defaults; simple linear relation")
@@ -27,6 +31,14 @@ equations = pysr(X, y,
                  unary_operators=[], binary_operators=["plus"],
                  niterations=10,
                  user_input=False)
-
 print(equations)
 assert equations.iloc[-1]['MSE'] < 1e-4
+
+print("Test 4 - text JAX export")
+x, y, z = sympy.symbols('x y z')
+cosx = 1.0 * sympy.cos(x) + y
+key = random.PRNGKey(0)
+X = random.normal(key, (1000, 2))
+true = 1.0 * jnp.cos(X[:, 0]) + X[:, 1]
+f, params = sympy2jax(cosx, [x])
+assert jnp.all(jnp.isclose(f(X, params), true)).item()