Allow user to pass extra torch operators

MilesCranmer · May 31, 2021 · 9269ce9 · 9269ce9
1 parent 4697288
commit 9269ce9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pysr/export_torch.py b/pysr/export_torch.py
@@ -160,7 +160,7 @@ def forward(self, X):
                 return self._node(symbols)
 
 
-def sympy2torch(expression, symbols_in):
+def sympy2torch(expression, symbols_in, extra_torch_mappings=None):
     """Returns a module for a given sympy expression with trainable parameters;
 
     This function will assume the input to the module is a matrix X, where
@@ -170,4 +170,4 @@ def sympy2torch(expression, symbols_in):
 
     _initialize_torch()
 
-    return SingleSymPyModule(expression, symbols_in)
+    return SingleSymPyModule(expression, symbols_in, extra_funcs=extra_torch_mappings)