Implement destructive in-place rewrites for Cholesky and Solve Ops

[jessegrabowski]

Description

All scipy.linalg functions offer an overwrite_a and/or overwrite_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

• [x] Closes #572
• [ ] Related to #

Checklist

• [x] Checked that the pre-commit linting/style checks pass
• [x] Included tests that prove the fix is effective or that the new feature works
• [ ] Added necessary documentation (docstrings and/or example notebooks)
• [x] If you are a pro: each commit corresponds to a relevant logical change

Type of change

• [x] New feature / enhancement
• [ ] Bug fix
• [ ] Documentation
• [ ] Maintenance
• [ ] Other (please specify):

[ricardoV94]

Can the Numba Cholesky make use of overwrite?

[jessegrabowski]

I think so? I added a check to copy the input matrix or not. I need to test it more carefully to make sure it does what I want it to do.

[ricardoV94]

Numba seems to be computing in integer? Anyway I guess this PR now depends on #578?

[jessegrabowski]

I had removed a .astype(input.dtype) to the output of Cholesky().perform, so it made the integer test fail. Everything should pass now. #578 is just a general numba speedup, the two shouldn't clash (I hope).

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 81.72043% with 17 lines in your changes missing coverage. Please review.

Project coverage is 81.73%. Comparing base (1c2bc8f) to head (cb0ec82). Report is 16 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/slinalg.py	77.96%	7 Missing and 6 partials :warning:
pytensor/tensor/blockwise.py	50.00%	1 Missing and 1 partial :warning:
pytensor/tensor/rewriting/blockwise.py	92.85%	1 Missing and 1 partial :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #577      +/-   ##
==========================================
- Coverage   81.73%   81.73%   -0.01%     
==========================================
  Files         183      183              
  Lines       47734    47814      +80     
  Branches    11611    11635      +24     
==========================================
+ Hits        39016    39080      +64     
- Misses       6523     6532       +9     
- Partials     2195     2202       +7     

Files with missing lines	Coverage Δ
pytensor/graph/op.py	88.02% <100.00%> (+0.12%)	:arrow_up:
pytensor/scan/rewriting.py	81.52% <ø> (ø)
pytensor/sparse/rewriting.py	76.15% <ø> (ø)
pytensor/tensor/random/rewriting/basic.py	93.19% <ø> (ø)
pytensor/tensor/rewriting/blas.py	88.21% <ø> (ø)
pytensor/tensor/rewriting/blas_scipy.py	82.35% <ø> (ø)
pytensor/tensor/rewriting/elemwise.py	90.84% <ø> (ø)
pytensor/tensor/rewriting/subtensor.py	90.12% <ø> (ø)
pytensor/typed_list/rewriting.py	100.00% <ø> (ø)
pytensor/tensor/blockwise.py	81.21% <50.00%> (-0.65%)	:arrow_down:
... and 2 more

... and 3 files with indirect coverage changes

[ricardoV94]

If you are a pro: each commit corresponds to a relevant logical change

You marked that, but the commits seem "dirty". Do you mind if I squash merge?

[jessegrabowski]

I don't mind at all, but I'd rather you explained it to me so I can do it right next time.

I read the link, but it suggested I write War and Peace in each commit message.

[ricardoV94]

I read the link, but it suggested I write War and Peace in each commit message.

Fair enough. The biggest point is each commit should be a self-contained logical change, so that you could in theory revert or checkout any of them and the codebase would still make sense as is. Usually stuff like fix typo, run-pre commit, add suggestions from review, intermediate changes that were ultimately not needed, disappear after "cleaning"

[jessegrabowski]

OK let me take a stab at cleaning it up

[ricardoV94]

OK let me take a stab at cleaning it up

Advice: backup the branch before you try :)

[jessegrabowski]

I just did an interactive rebase and squashed things together. Is that all I needed to do?

[bwengals]

sorry for the delay, ya'll are fast! Taking a look now

[bwengals]

Works great for me (running the test you wrote)! Are you planning on doing all of them in this PR? the solves, eig's, etc?