Fix feature selection for JAX export

pysr/export_jax.py

@@ -109,7 +109,7 @@ def _initialize_jax():
         jsp = _jsp
 
 
-def sympy2jax(expression, symbols_in, extra_jax_mappings=None):
+def sympy2jax(expression, symbols_in, selection=None, 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,6 +192,9 @@ def sympy2jax(expression, symbols_in, 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 = {}

pysr/sr.py

@@ -1740,6 +1740,7 @@ class PySRRegressor(BaseEstimator, RegressorMixin, MultiOutputMixin):
                     func, params = sympy2jax(
                         eqn,
                         sympy_symbols,
+                        selection=self.selection_mask_,
                         extra_jax_mappings=self.extra_jax_mappings,
                     )
                     jax_format.append({"callable": func, "parameters": params})