Propagate and check torch/jax mappings

pysr/export_jax.py

@@ -51,7 +51,7 @@ _jnp_func_lookup = {
 }
 
 
-def sympy2jaxtext(expr, parameters, symbols_in):
+def sympy2jaxtext(expr, parameters, symbols_in, extra_jax_mappings=None):
     if issubclass(expr.func, sympy.Float):
         parameters.append(float(expr))
         return f"parameters[{len(parameters) - 1}]"
@@ -61,8 +61,15 @@ def sympy2jaxtext(expr, parameters, symbols_in):
         return (
             f"X[:, {[i for i in range(len(symbols_in)) if symbols_in[i] == expr][0]}]"
         )
-    _func = _jnp_func_lookup[expr.func]
-    args = [sympy2jaxtext(arg, parameters, symbols_in) for arg in expr.args]
+    if extra_jax_mappings is None:
+        extra_jax_mappings = {}
+    _func = {**_jnp_func_lookup, **extra_jax_mappings}[expr.func]
+    args = [
+        sympy2jaxtext(
+            arg, parameters, symbols_in, extra_jax_mappings=extra_jax_mappings
+        )
+        for arg in expr.args
+    ]
     if _func == MUL:
         return " * ".join(["(" + arg + ")" for arg in args])
     if _func == ADD:
@@ -92,7 +99,7 @@ def _initialize_jax():
         jsp = _jsp
 
 
-def sympy2jax(expression, symbols_in, selection=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)
@@ -170,7 +177,9 @@ def sympy2jax(expression, symbols_in, selection=None):
     global jsp
 
     parameters = []
-    functional_form_text = sympy2jaxtext(expression, parameters, symbols_in)
+    functional_form_text = sympy2jaxtext(
+        expression, parameters, symbols_in, extra_jax_mappings
+    )
     hash_string = "A_" + str(abs(hash(str(expression) + str(symbols_in))))
     text = f"def {hash_string}(X, parameters):\n"
     if selection is not None:

pysr/sr.py

@@ -289,6 +289,20 @@ def pysr(
     if len(variable_names) == 0:
         variable_names = [f"x{i}" for i in range(X.shape[1])]
 
+    if extra_jax_mappings is not None:
+        for key, value in extra_jax_mappings:
+            if not isinstance(value, str):
+                raise NotImplementedError(
+                    "extra_jax_mappings must have keys that are strings! e.g., {sympy.sqrt: 'jnp.sqrt'}."
+                )
+
+    if extra_torch_mappings is not None:
+        for key, value in extra_jax_mappings:
+            if not callable(value):
+                raise NotImplementedError(
+                    "extra_torch_mappings must be callable functions! e.g., {sympy.sqrt: torch.sqrt}."
+                )
+
     use_custom_variable_names = len(variable_names) != 0
 
     _check_assertions(
@@ -996,14 +1010,24 @@ def get_hof(
             if output_jax_format:
                 from .export_jax import sympy2jax
 
-                func, params = sympy2jax(eqn, sympy_symbols, selection)
+                func, params = sympy2jax(
+                    eqn,
+                    sympy_symbols,
+                    selection=selection,
+                    extra_jax_mappings=extra_jax_mappings,
+                )
                 jax_format.append({"callable": func, "parameters": params})
 
             # Torch:
             if output_torch_format:
                 from .export_torch import sympy2torch
 
-                module = sympy2torch(eqn, sympy_symbols, selection=selection)
+                module = sympy2torch(
+                    eqn,
+                    sympy_symbols,
+                    selection=selection,
+                    extra_torch_mappings=extra_torch_mappings,
+                )
                 torch_format.append(module)
 
             curMSE = output.loc[i, "MSE"]