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

MilesCranmer / PySR

High-Performance Symbolic Regression in Python and Julia

https://astroautomata.com/PySR

Apache License 2.0

2.09k stars 197 forks source link

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

Open MilesCranmer opened 9 months ago

MilesCranmer commented 9 months ago

Thanks to @j-thib for pointing this out, I didn't realize the SymPy maps were not built-in.

However this gets fairly complicated as we need to map sympy.Piecewise into torch/jax code.

TODO:

[x] Add sympy2torch for Piecewise
[ ] Add sympy2jax for Piecewise
[ ] Add tests for all new branches. We should reach 100% coverage on the diff because the code gets fairly convoluted.

MilesCranmer commented 1 month ago

If you go to the files tab, you can click the + next to a line of code where you want to add it. Then click "suggest". Then I can add your code as a commit on the PR. If you can't figure it out no worries!