Add Cholesky dot(L, L.T) rewrite and remove sandbox module

[dehorsley]

Motivation for these changes

As discussed in pymc-devs/pymc#6717, it is useful in PyMC to rewrite cholesky(dot(L,L.T)) -> L for use in MvNormal logp.

The equivalent was included for the upper triangular factor case.

Implementation details

This uses a lower_triangular or upper_triangular tag to ensure the input matrix is the correct form.

New features

• Add rewrite for cholesky(dot(L, L.T))

[ricardoV94]

Thanks @dehorsley.

Note that PyMC will still have a messed up dlogp, so it may not be enough to fix performance, but it should at least help.
It's much more useful for general PyTensor graphs, so thanks for taking a look!

[codecov-commenter]

Codecov Report

Merging #303 (d14601d) into main (6b43b43) will decrease coverage by 0.01%.
The diff coverage is 61.11%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #303      +/-   ##
==========================================
- Coverage   80.50%   80.49%   -0.01%     
==========================================
  Files         154      154              
  Lines       45186    45204      +18     
  Branches    11046    11056      +10     
==========================================
+ Hits        36376    36387      +11     
- Misses       6608     6611       +3     
- Partials     2202     2206       +4     

Impacted Files	Coverage Δ
pytensor/sandbox/linalg/ops.py	71.07% <61.11%> (-1.75%)	⬇️

[dehorsley]

@ricardoV94 one question about this: it's not clear to me that when pytensor/sandbox/linalg/ops.py is imported. While testing this in pymc, the optimisations aren't registered up unless I explicitly import the module. Should this change be put in a different module?

[ricardoV94]

@ricardoV94 one question about this: it's not clear to me that when pytensor/sandbox/linalg/ops.py is imported. While testing this in pymc, the optimisations aren't registered up unless I explicitly import the module. Should this change be put in a different module?

We should move the rewrites (also the old ones unless you see any potential problems with them) to tensor/rewriting/linalg.py (new file). Then it should be imported in rewriting/__init__.py so they are registered by default.

I think we can just delete the whole sandbox module at this point. It only contains a dummy Minimal(Op) and this spectral_radius_bound that I think we can get rid of. I doubt anyone is using it since we don't advertise it.

Do you mind doing that?

[dehorsley]

We should move the rewrites (also the old ones unless you see any potential problems with them) to tensor/rewriting/linalg.py (new file). Then it should be imported in rewriting/__init__.py so they are registered by default.

I've hit a circular dependency with this:

tensor/blas.py has

from pytensor.tensor.rewriting.elemwise import local_dimshuffle_lift

And tensor/rewriting/linalg.py

from pytensor.tensor.blas import Dot22

Maybe the "right" thing to do is to move the rewrites in tensor/blas.py to tensor/rewriting/blas.py. However, the relevant rewrites in rewriting/linalg.py in the stabilize group and run before the blas specialization, so there shouldn't be any Dot22s in the graph at that point. Any thoughts?

EDIT: I've gone with the latter.

[ricardoV94]

Looks good except for the rewrite change that's still there.

[dehorsley]

Thanks, done.

By the way, I got a lot of mypy errors in unrelated files as part of the pre-commit script, is that anything to be concerned about?

[ricardoV94]

Sometimes pre-commit fails locally because it's run only on the changed files but passes here where it is run globally. I'll make the tests run now and we can see

[ricardoV94]

Is this the failure you saw? https://github.com/pymc-devs/pytensor/actions/runs/5055814502/jobs/9072710031?pr=303#step:4:434

I guess it's just a case of removing it?

[dehorsley]

Yup that's what I saw locally too. Any idea what to do?

[ricardoV94]

Yup that's what I saw locally too. Any idea what to do?

I think we can remove those # type: ignore comments, but you might need to disable your pre-commit to be able to commit that change ??

@michaelosthege any guidance here?

[dehorsley]

Yeah, I had to disable pre-commit to commit the change. Should otherwise be good to go if the tests run fine. They ran OK locally so 🤞.

[michaelosthege]

error: Unused "type: ignore" comment

sounds pretty much like they can be removed

[dehorsley]

looks like the fixed it locally

[dehorsley]

The failing tests seem to be to do with some numba random issues unrelated to this PR. Are you happy otherwise @ricardoV94 ?

[ricardoV94]

@dehorsley the numba tests were fixed in another PR.

I rebased from main and also rearranged the commits so that the library is "coherent" after each step.
Careful with pulling from your remote if you need to do more changes.

If tests pass I think this is good to go.
Let me know if I messed up something or if there is something else you would like to change.

[ricardoV94]

Thanks @dehorsley

Do you want to investigate how much this helps on the PyMC side? We will have to add the tag to the value variable when it comes out of that transform