Add test for pytorch export

.github/workflows/CI.yml

@@ -61,17 +61,23 @@ jobs:
             python setup.py install
       - name: "Install Coverage tool"
         run: pip install coverage coveralls
+      - name: "Run tests"
+        run: coverage run --source=pysr --omit='*/feynman_problems.py' -m unittest test.test
+        shell: bash
       - name: "Install JAX"
         if: matrix.os != 'windows-latest'
         run: pip install jax jaxlib # (optional import)
         shell: bash
-      - name: "Run tests"
-        run: coverage run --source=pysr --omit='*/feynman_problems.py' -m unittest test.test
-        shell: bash
       - name: "Run JAX tests"
         if: matrix.os != 'windows-latest'
         run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_jax
         shell: bash
+      - name: "Install Torch"
+        run: pip install torch # (optional import)
+        shell: bash
+      - name: "Run Torch tests"
+        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_torch
+        shell: bash
       - name: Coveralls
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

test/test_torch.py

@@ -0,0 +1,12 @@
+import numpy as np
+from pysr import sympy2torch
+import torch
+import sympy
+
+print("Test Torch 1 - test export")
+x, y, z = sympy.symbols('x y z')
+cosx = 1.0 * sympy.cos(x) + y
+X = torch.randn((1000, 2))
+true = 1.0 * torch.cos(X[:, 0]) + X[:, 1]
+torch_module = sympy2torch(cosx, [x, y, z])
+assert jnp.all(jnp.isclose(torch_module(X), true)).item()