Only test jax on non-Windows systems

.github/workflows/CI.yml

@@ -58,9 +58,15 @@ jobs:
         run: |
             python -m pip install --upgrade pip
             pip install -r requirements.txt
-            pip install jax jaxlib # (optional import)
             python setup.py install
+      - name: "Install JAX"
+        if: matrix.os != 'windows-latest'
+        run: pip install jax jaxlib # (optional import)
         shell: bash
       - name: "Run tests"
         run: python test/test.py
         shell: bash
+      - name: "Run JAX tests"
+        if: matrix.os != 'windows-latest'
+        run: python test/test_jax.py
+        shell: bash

test/test.py

@@ -1,8 +1,5 @@
 import numpy as np
-from pysr import pysr, sympy2jax
-from jax import numpy as jnp
-from jax import random
-from jax import grad
+from pysr import pysr
 import sympy
 X = np.random.randn(100, 5)
 
@@ -33,12 +30,3 @@ equations = pysr(X, y,
                  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()

test/test_jax.py

@@ -0,0 +1,15 @@
+import numpy as np
+from pysr import pysr, sympy2jax
+from jax import numpy as jnp
+from jax import random
+from jax import grad
+import sympy
+
+print("Test JAX 1 - test 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()