Add missing PyTorch/JAX export for logical_or, logical_and, and relu

pysr/export_jax.py

@@ -69,7 +69,7 @@ def sympy2jaxtext(expr, parameters, symbols_in, extra_jax_mappings=None):
  _func = {**_jnp_func_lookup, **extra_jax_mappings}[expr.func]
  except KeyError:
  raise KeyError(
- f"Function {expr.func} was not found in JAX function mappings."
+ f"Function {expr.func} was not found in JAX function mappings. "
  "Please add it to extra_jax_mappings in the format, e.g., "
  "{sympy.sqrt: 'jnp.sqrt'}."
  )

pysr/export_torch.py

@@ -33,6 +33,70 @@ def _initialize_torch():
 
  torch = _torch
 
+ # Allows PyTorch to map Piecewise functions:
+ def expr_cond_pair(expr, cond):
+ if isinstance(cond, torch.Tensor) and not isinstance(expr, torch.Tensor):
+ expr = torch.tensor(expr, dtype=cond.dtype, device=cond.device)
+ elif isinstance(expr, torch.Tensor) and not isinstance(cond, torch.Tensor):
+ cond = torch.tensor(cond, dtype=expr.dtype, device=expr.device)
+ else:
+ return expr, cond
+
+ # First, make sure expr and cond are same size:
+ if expr.shape != cond.shape:
+ if len(expr.shape) == 0:
+ expr = expr.expand(cond.shape)
+ elif len(cond.shape) == 0:
+ cond = cond.expand(expr.shape)
+ else:
+ raise ValueError(
+ "expr and cond must have same shape, or one must be a scalar."
+ )
+ return expr, cond
+
+ def if_then_else(*conds):
+ a, b, c = conds
+ return torch.where(
+ a, torch.where(b, True, False), torch.where(c, True, False)
+ )
+
+ def piecewise(*expr_conds):
+ output = None
+ already_used = None
+ for expr, cond in expr_conds:
+ if not isinstance(cond, torch.Tensor) and not isinstance(
+ expr, torch.Tensor
+ ):
+ # When we just have scalars, have to do this a bit more complicated
+ # due to the fact that we need to evaluate on the correct device.
+ if output is None:
+ already_used = cond
+ output = expr if cond else 0.0
+ else:
+ if not isinstance(output, torch.Tensor):
+ output += expr if cond and not already_used else 0.0
+ already_used = already_used or cond
+ else:
+ expr = torch.tensor(
+ expr, dtype=output.dtype, device=output.device
+ ).expand(output.shape)
+ output += torch.where(
+ cond & ~already_used, expr, torch.zeros_like(expr)
+ )
+ already_used = already_used | cond
+ else:
+ if output is None:
+ already_used = cond
+ output = torch.where(cond, expr, torch.zeros_like(expr))
+ else:
+ output += torch.where(
+ cond.bool() & ~already_used, expr, torch.zeros_like(expr)
+ )
+ already_used = already_used | cond.bool()
+ return output
+
+ # TODO: Add test that makes sure tensors are on the same device
+
  _global_func_lookup = {
  sympy.Mul: _reduce(torch.mul),
  sympy.Add: _reduce(torch.add),
@@ -81,6 +145,12 @@ def _initialize_torch():
  sympy.Heaviside: torch.heaviside,
  sympy.core.numbers.Half: (lambda: 0.5),
  sympy.core.numbers.One: (lambda: 1.0),
+ sympy.logic.boolalg.Boolean: lambda x: x,
+ sympy.logic.boolalg.BooleanTrue: (lambda: True),
+ sympy.logic.boolalg.BooleanFalse: (lambda: False),
+ sympy.functions.elementary.piecewise.ExprCondPair: expr_cond_pair,
+ sympy.Piecewise: piecewise,
+ sympy.logic.boolalg.ITE: if_then_else,
  }
 
  class _Node(torch.nn.Module):
@@ -125,7 +195,7 @@ def __init__(self, *, expr, _memodict, _func_lookup, **kwargs):
  self._torch_func = _func_lookup[expr.func]
  except KeyError:
  raise KeyError(
- f"Function {expr.func} was not found in Torch function mappings."
+ f"Function {expr.func} was not found in Torch function mappings. "
  "Please add it to extra_torch_mappings in the format, e.g., "
  "{sympy.sqrt: torch.sqrt}."
  )
@@ -153,7 +223,13 @@ def forward(self, memodict):
  arg_ = arg(memodict)
  memodict[arg] = arg_
  args.append(arg_)
- return self._torch_func(*args)
+ try:
+ return self._torch_func(*args)
+ except Exception as err:
+ # Add information about the current node to the error:
+ raise type(err)(
+ f"Error occurred in node {self._sympy_func} with args {args}"
+ )
 
  class _SingleSymPyModule(torch.nn.Module):
  """SympyTorch code from https://github.com/patrick-kidger/sympytorch"""