Fix numba implementation of cholesky not setting off-diag entries to zero

[aseyboldt]

Description

The cholesky blas function returns an array with values stored in the triagonal part of the array that is zero by definition.

This is the equivalent of the scipy code here: https://github.com/scipy/scipy/blob/d5e1e03ac12fb5ad7fa4949e7ffea2137e780b15/scipy/linalg/flapack_pos_def.pyf.src#L132

fixes #814

Checklist

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

Type of change

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

[ricardoV94]

Any reason why not to extend/reuse the test below?

[aseyboldt]

No good one, fixed.

[jessegrabowski]

Nice, good catch.

#814 was about originally about gradients, so perhaps we should test them explicitly?

[ricardoV94]

Nice, good catch.

#814 was about originally about gradients, so perhaps we should test them explicitly?

I would say no, as long as someone checked the motivating example and it seems to work. The cases where bugs are found can be much more complex/expensive than the final regression tests?

[aseyboldt]

I checked the original example.
I think the reason the problem only showed up in the gradient is that in the forward pass the values that should be zero were not used. The triagonal solve function just ignores what's in there, but I think in the backward pass we had a matrix mult or so which does not ignore those values.

[ricardoV94]

Isn't compare_numba_and_py already checking numba and python implementations match?

[aseyboldt]

It should, but didn't find the problem because compare_numba_and_py calls the functions without fast_run rewrites. That excluded Blockwise, and because that doesn't have a numba implementation it ended up calling the python code for the numba and the python version.

I included the rewrite that remove useless Blockwise ops into the set of rewrites we do in testing.
Maybe we should actually just use the default numba mode?

[ricardoV94]

Ahh... okay so hopefully just a temporary blind spot. Let's keep it like you did now

[codecov]

Codecov Report

Attention: Patch coverage is 12.50000% with 7 lines in your changes missing coverage. Please review.

Project coverage is 80.86%. Comparing base (ee0075b) to head (b766b21).
Report is 168 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/numba/dispatch/slinalg.py	12.50%	7 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #816      +/-   ##
==========================================
+ Coverage   80.84%   80.86%   +0.01%     
==========================================
  Files         163      163              
  Lines       46890    46897       +7     
  Branches    11468    11473       +5     
==========================================
+ Hits        37909    37923      +14     
+ Misses       6768     6757      -11     
- Partials     2213     2217       +4     

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/slinalg.py	52.98% <12.50%> (+5.44%)	⬆️

... and 1 file with indirect coverage changes