Implement gradient for SVD

[jessegrabowski]

Implement gradients for SVD

Description

I just copied the implementation from autograd here. The base case compute_uv=False, full_matrices=False works and test passes.

Marked as a draft because I'm not sure how to handle multiple output gradients in a test. I'm not sure my copy/paste job is correct until I verify (there are some slight differences between our SVD and the reference I was following -- I think we return V, but they returned V.conj().T. This is a numpy vs matlab thing I think.

Even if we compute all the matrices, we only get a gradient for the singular values. Is there a way to return [NoGrad, GradientGraph, NoGrad]?

Related Issue

• Closes Implement gradient of SVD #56
• 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

All modified and coverable lines are covered by tests ✅

Project coverage is 80.81%. Comparing base (eb18f0e) to head (196b5e4).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #614      +/-   ##
==========================================
+ Coverage   80.78%   80.81%   +0.02%     
==========================================
  Files         162      162              
  Lines       46757    46800      +43     
  Branches    11440    11449       +9     
==========================================
+ Hits        37773    37820      +47     
+ Misses       6735     6733       -2     
+ Partials     2249     2247       -2     

Files	Coverage Δ
pytensor/tensor/nlinalg.py	95.22% <100.00%> (+0.46%)	⬆️

... and 2 files with indirect coverage changes

[ricardoV94]

Is there a way to return [NoGrad, GradientGraph, NoGrad]?

Yes, but what's the reason for NoGrad? See different flags in section: https://pytensor.readthedocs.io/en/latest/extending/op.html#gradient

[ricardoV94]

For fun reference: https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf

[jessegrabowski]

I'm very confident the gradients are implemented correctly, but I'm getting shape errors in the gradient check for the non-square case. As usual, there's also no complex support -- I didn't even bother to add the xfail tests this time, but I could.

I am hoping to get a 2nd pair of eyes on it to see if I'm missing something obvious. I wonder if it's related to the use of shared variables in verify_grad? There's no need for it -- we should just pass the random project matrices as arguments to the compiled function.

[ricardoV94]

I wonder if it's related to the use of shared variables in verify_grad? There's no need for it -- we should just pass the random project matrices as arguments to the compiled function.

That shouldn't be a problem. I'll try to have a look, but that utility take a while to debug.

[jessegrabowski]

False alarm, I had a dumb mistake in the gradients. This should be good to go now.

I still hate shared variables.