Test selection inside jax/torch

test/test_jax.py

@@ -17,9 +17,9 @@ class TestJAX(unittest.TestCase):
         f, params = sympy2jax(cosx, [x, y, z])
         self.assertTrue(jnp.all(jnp.isclose(f(X, params), true)).item())
     def test_pipeline(self):
-        X = np.random.randn(100, 2)
+        X = np.random.randn(100, 10)
         equations = pd.DataFrame({
-            'Equation': ['1.0', 'cos(x0)', 'square(cos(x0))'],
+            'Equation': ['1.0', 'cos(x1)', 'square(cos(x1))'],
             'MSE': [1.0, 0.1, 1e-5],
             'Complexity': [1, 2, 3]
             })
@@ -30,7 +30,7 @@ class TestJAX(unittest.TestCase):
         equations = get_hof(
                 'equation_file.csv', n_features=2, variables_names='x0 x1'.split(' '),
                 extra_sympy_mappings={}, output_jax_format=True,
-                multioutput=False, nout=1)
+                multioutput=False, nout=1, selection=[1, 2, 3])
 
         jformat = equations.iloc[-1].jax_format
         np.testing.assert_almost_equal(

test/test_torch.py

@@ -16,9 +16,9 @@ class TestTorch(unittest.TestCase):
                 np.all(np.isclose(torch_module(X).detach().numpy(), true.detach().numpy()))
         )
     def test_pipeline(self):
-        X = np.random.randn(100, 2)
+        X = np.random.randn(100, 10)
         equations = pd.DataFrame({
-            'Equation': ['1.0', 'cos(x0)', 'square(cos(x0))'],
+            'Equation': ['1.0', 'cos(x1)', 'square(cos(x1))'],
             'MSE': [1.0, 0.1, 1e-5],
             'Complexity': [1, 2, 3]
             })
@@ -29,7 +29,7 @@ class TestTorch(unittest.TestCase):
         equations = get_hof(
                 'equation_file.csv', n_features=2, variables_names='x0 x1'.split(' '),
                 extra_sympy_mappings={}, output_torch_format=True,
-                multioutput=False, nout=1)
+                multioutput=False, nout=1, selection=[1, 2, 3])
 
         tformat = equations.iloc[-1].torch_format
         np.testing.assert_almost_equal(