-
Notifications
You must be signed in to change notification settings - Fork 104
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
Implement destructive in-place rewrites for Cholesky and Solve Ops #1028
Conversation
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1028 +/- ##
==========================================
- Coverage 81.89% 81.88% -0.01%
==========================================
Files 182 182
Lines 47778 47858 +80
Branches 8597 8618 +21
==========================================
+ Hits 39126 39190 +64
- Misses 6487 6494 +7
- Partials 2165 2174 +9
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Very cool, glad to see this over the finish line (with a bunch of extended functionality!)
I left mostly nitpicks. This could be merged as-is without any changes IMO.
My biggest criticism of the approach, if I had to offer one, would be the use of a list of integers for allowed_inplace_inputs
, which you seem to be aware of yoruself because of the extended comment in slinalg.py:L521-523
. I don't have a better alternative (what came immediately to mind is to always pass a list of bool of length len(self.inputs)
, and check the appropriate positions for True. but it's no better). I guess it's frustrating that we sometimes encode informaton about codes with dicts (destroy_map
) and other times with lists (allowed_inplace_inputs
). It's somewhat difficult as a dev to know what to expect.
Co-authored-by: Ricardo Vieira <[email protected]>
1de1726
to
15384ef
Compare
Description
Reviving accidentally closed #577
All scipy.linalg functions offer an
overwrite_a
and/oroverwrite_b
argument that can enhance performance by re-using the input memory for the outputs. This PR implements re-writes that will set these flags to True at compile time.These rewrites are also a nice example for the in-place docs here, so I'll update them with an example as a later commit to this PR.
We added a general inplace functionality to Blockwise, that automatically makes the core_op inplace. This will translate to the core_op if the Blockwise is later removed due to being useless (no batch dims actually exist).
Related Issue
Op
#572Checklist
Type of change
📚 Documentation preview 📚: https://pytensor--1028.org.readthedocs.build/en/1028/