Refactor nlinalg.norm to match np.linalg.norm signature and functionaly

[jessegrabowski]

Description

Currently, pt.linalg.norm does not match the signature of np.linalg.norm -- it is missing the axis and keepdims arguments, and has no logic for handling complex-valued matrices. This PR fixes this by copy-pasting the np.linalg.norm code to pytensor.

In addition, some extra logic is needed to allow for tensor-valued inputs to norm. This part is still a WIP. The tests need to be refactored to test all the new combinations of keyword arguments as well.

This is a preliminary step to addressing pymc-devs/pymc#7101, because the correct transformation between unconstrained lower-triangular matrices to cholesky factorized matrices involves row-normalizing the unconstrained lower-triangular matrix; see here.

Related Issue

• Closes #
• Related to BUG: LKJCorr breaks when used as covariance with MvNormal pymc#7101

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-commenter]

Codecov Report

Attention: 16 lines in your changes are missing coverage. Please review.

Comparison is base (ee7c946) 80.83% compared to head (2812aca) 80.83%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #588      +/-   ##
==========================================
- Coverage   80.83%   80.83%   -0.01%     
==========================================
  Files         162      162              
  Lines       46743    46797      +54     
  Branches    11417    11434      +17     
==========================================
+ Hits        37786    37827      +41     
- Misses       6714     6722       +8     
- Partials     2243     2248       +5     

Files	Coverage Δ
pytensor/tensor/nlinalg.py	94.51% <81.17%> (-3.58%)	⬇️

... and 2 files with indirect coverage changes

[jessegrabowski]

I ended up just manually handing the broadcasting cases. Curious if my original approach could work somehow, because the solution I converged on duplicates a lot of the logic handled automatically by Blockwise.

I also had to make SVD Blockwise en passant, so that's a nice bonus from this PR.

[ricardoV94]

These changes look amazing!

Left some comments, but I don't see anything that should be done structurally differently

[jessegrabowski]

I reorganized this PR into two separate commits, one for SVD and one for norm. It should be good to go now.

[ricardoV94]

Did we have any rewrites targeting SVD Ops? Now they we Blockwise they would need to be tweaked, but since tests pass I assume we don't have any

[jessegrabowski]

No, it was kind of an orphan Op, with no rewrites or gradients.