Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 23 commits into
base: master
Choose a base branch
from

Conversation

MilesCranmer
Copy link
Owner

@MilesCranmer MilesCranmer commented Sep 19, 2023

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:

  • 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 MilesCranmer changed the title Add missing SymPy maps for logical_or, logical_and, and relu Add missing PyTorch/JAX export for logical_or, logical_and, and relu Dec 14, 2023
Copy link

@tbuckworth tbuckworth left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In order to avoid the following error:
KeyError: 'Function ITE was not found in Torch function mappings.Please add it to extra_torch_mappings in the format, e.g., {sympy.sqrt: torch.sqrt}.'

This function should be added to the mappings:

def if_then_else(*conds):
    a, b, c = conds
    return torch.where(a, torch.where(b, True, False), torch.where(c, True, False))

extra_torch_mappings = {sympy.logic.boolalg.ITE: if_then_else}

Additionally, lines 87-89 can be replaced with:

output += torch.where(
                    cond.bool() & ~already_used, expr, torch.zeros_like(expr)
                )
                already_used = already_used | cond.bool()

as when cond is a float (1.) the code fails without cond.bool()

@MilesCranmer
Copy link
Owner Author

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!

Copy link

@tbuckworth tbuckworth left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added if_then_else as ITE mapping.

Changed cond to cond.bool() to avoid error when cond it Float (1.).

pysr/export_torch.py Show resolved Hide resolved
pysr/export_torch.py Show resolved Hide resolved
pysr/export_torch.py Outdated Show resolved Hide resolved
pysr/export_torch.py Outdated Show resolved Hide resolved
pysr/export_torch.py Outdated Show resolved Hide resolved
@MilesCranmer MilesCranmer changed the base branch from store-options to master June 3, 2024 16:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants