Adds functions to rewrite cholesky decomposition of identity and diagonal matrices

[tanish1729]

Description

Add a function to rewrite cholesky(eye) -> eye and cholesky(diag) -> sqrt(diag)

Related Issue

• Related to ENH: Implement COLA library rewrites for linear algebra functions #573

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):

[codecov]

Codecov Report

Attention: Patch coverage is 93.93939% with 2 lines in your changes missing coverage. Please review.

Project coverage is 81.76%. Comparing base (3e98b9f) to head (cf87362).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/rewriting/linalg.py	93.93%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #925   +/-   ##
=======================================
  Coverage   81.75%   81.76%           
=======================================
  Files         183      183           
  Lines       47763    47796   +33     
  Branches    11619    11635   +16     
=======================================
+ Hits        39050    39081   +31     
+ Misses       6523     6522    -1     
- Partials     2190     2193    +3     

Files with missing lines	Coverage Δ
pytensor/tensor/rewriting/linalg.py	90.93% <93.93%> (+0.29%)	⬆️

... and 2 files with indirect coverage changes

[tanish1729]

@jessegrabowski @ricardoV94

[jessegrabowski]

Some comments and suggestions. I'd also ask that you add a test to make sure the cases that should not be rewritten are, in fact, not rewritten. For example, pt.linalg.cholesky(pt.eye(7, k=-1) * x)

[jessegrabowski]

Looking great, just a couple nitpicks. Could you also add a batch case to each of test_cholesky_diag_from_diag and test_cholesky_diag_from_eye_mul, then I think we'll be there.

[jessegrabowski]

LGTM