Skip to content

Cleanup fusion rewrite database #369

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

Merged
merged 3 commits into from
Jul 4, 2023

Conversation

ricardoV94
Copy link
Member

@ricardoV94 ricardoV94 commented Jul 3, 2023

  • Fix a couple of bugs in the careduce_fusion rewrite. These are very much related to some discussion on
    Add rewrite for Sum(MakeVector) #346 about tricky aspects of the CAReduce Ops. This rewrite was accidentally rarely used because it was registered outside the fusion database, and ran before the fusion rewrites. This explains why the failures where hidden for so long. It is now included in the same database, which is also related to the next point:
  • Remove the flag to create the fusion database. This unnecessarily complicated the codebase logic, and the flag was not even "mutable". The same behavior can be done by excluding "fusion" tags from the compilation mode.
  • Remove specialize_device database. This seems to be a left-over from the time Theano supported GPU. Only one rewrite (of dubious value) existed in it, which I have moved into the standard specialize database.

Spinoff from #361

Same behavior can be obtained with `optimizer_excluding`

The `local_careduce_rewrite` is now included in this database. Otherwise it would usually not be applied because it ran before the fusion rewrites
@ricardoV94 ricardoV94 changed the title Don't run inplace and fusion rewrites in JAX Cleanup fusion rewrite database Jul 3, 2023
@ricardoV94 ricardoV94 requested a review from michaelosthege July 3, 2023 19:53
Comment on lines +1186 to +1192
pytest.param(
at_all,
np.all,
marks=pytest.mark.xfail(
reason="Rewrite logic does not support all CAReduce"
),
),
Copy link
Member

Choose a reason for hiding this comment

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

TIL !

@ricardoV94 ricardoV94 merged commit e20dd0b into pymc-devs:main Jul 4, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants