Fixed jax export compatibility with refactor

pysr/export_jax.py

@@ -109,7 +109,7 @@ def _initialize_jax():
         jsp = _jsp
 
 
-def sympy2jax(expression, symbols_in, selection=None, extra_jax_mappings=None):
+def sympy2jax(expression, symbols_in, extra_jax_mappings=None):
     """Returns a function f and its parameters;
     the function takes an input matrix, and a list of arguments:
             f(X, parameters)
@@ -192,9 +192,6 @@ def sympy2jax(expression, symbols_in, selection=None, extra_jax_mappings=None):
     )
     hash_string = "A_" + str(abs(hash(str(expression) + str(symbols_in))))
     text = f"def {hash_string}(X, parameters):\n"
-    if selection is not None:
-        # Impose the feature selection:
-        text += f"    X = X[:, {list(selection)}]\n"
     text += "    return "
     text += functional_form_text
     ldict = {}