Implement gradient for QR decomposition #1303

educhesne · 2025-03-19T10:38:51Z

Description

Implement the support for gradient in the QR decomposition. It is only defined when R is a square matrix.

Related Issue

Closes Implement gradient for QR decomposition #1099

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

📚 Documentation preview 📚: https://pytensor--1303.org.readthedocs.build/en/1303/

jessegrabowski

This is really great! The most important change that we need here is to use solve_triangular instead of inverse.

If you're willing, I would also like to move the QR function from nlinalg to slinalg, and use the scipy signature in the perform method. If you're not that's fine, and we can do that in a future PR.

pytensor/tensor/nlinalg.py

jessegrabowski · 2025-03-19T11:01:18Z

pytensor/tensor/nlinalg.py

+            Rinvt = _H(inv(R))
+            A_bar = Q @ ((ptb.tril(R_dRt - _H(R_dRt), k=-1)) @ Rinvt + dR)


Use a solve here instead of explicitly inverting R. You can use solve_triangular to exploit the structure of R, and set trans=1 instead of inverting R at all.

pytensor/tensor/nlinalg.py

tests/tensor/test_nlinalg.py

educhesne · 2025-03-19T13:05:00Z

Thanks for the quick review !
I cannot import solve_triangular in pt.nlinalg (it creates circular dependencies), so I'll move qr from nlinalg to slinalg directly.
The scipy signature is different, in particular mode values change (from complete, reduced to full, economic) and the default mode is complete/full rather than reduced/economic; should qr reflects the scipy signature ? or keep the numpy one ?

jessegrabowski · 2025-03-19T13:07:11Z

to import it without circular, use from pytensor.tensor.slinalg import solve_triangular inside the gradient function

Forget I brought up changing the signature, this PR is already amazing so let's just focus on gradients here.

codecov · 2025-03-19T13:30:32Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.00%. Comparing base (a149f6c) to head (4edc698).
Report is 12 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1303      +/-   ##
==========================================
- Coverage   82.00%   82.00%   -0.01%     
==========================================
  Files         188      188              
  Lines       48478    48530      +52     
  Branches     8665     8678      +13     
==========================================
+ Hits        39755    39797      +42     
- Misses       6575     6582       +7     
- Partials     2148     2151       +3

Files with missing lines	Coverage Δ
pytensor/tensor/nlinalg.py	`95.63% <100.00%> (+0.36%)`	⬆️

... and 6 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

educhesne · 2025-03-19T14:29:04Z

ok thanks ! I didn't use the trans argument of solve_triangular though; would that simplify the expression ? (it works anyway)

jessegrabowski · 2025-03-19T17:37:39Z

No worries about the trans argument, I think I misunderstood what was needed.

This is really great! Thanks for your contribution. It looks nice to me. Going to ask @ricardoV94 for a final sign off then merge.

ricardoV94

This is awesome!

Just some small change requests

ricardoV94 · 2025-03-21T10:56:00Z

pytensor/tensor/nlinalg.py

+        if self.mode == "raw" or (self.mode == "complete" and m != n):
+            raise NotImplementedError("Gradient of qr not implemented")
+
+        elif m < n:


gradient should work for non-static shapes, A = matrix("A", shape=(None, None)). If you need the shape check add it symbolically with assert_op

I replaced the shape checks with Assert on the output shape. Let me know if that's what you had in mind, I'm not very familiar with symbolic shapes

ricardoV94 · 2025-03-21T10:56:18Z

pytensor/tensor/nlinalg.py

+            return ptb.triu(x, k=0) + _H(ptb.triu(x, k=1))
+
+        if self.mode == "raw" or (self.mode == "complete" and m != n):
+            raise NotImplementedError("Gradient of qr not implemented")


Mention what exactly makes it non-implemented

* for mode=reduced or mode=r, all input shapes are accepted * for mode=complete, shapes m x n where m <= n are accepted

pytensor/tensor/nlinalg.py

jessegrabowski

lgtm

educhesne · 2025-03-26T11:02:30Z

I'm really sorry, wait a bit before merging, I may have mixed conjugate transpose and transpose. The tests are only done with real matrix, so that wouldn't be caught.

jessegrabowski · 2025-03-26T11:23:49Z

Can you try adding a complex test? You can follow the template here. Those tests don't work because numba is fussy. I think this case won't work either because we don''t support complex gradients, but you can still add the tests and xfail them (it's aspirational that way!)

educhesne · 2025-03-26T11:47:31Z

Ok, done.
I've rechecked the conjugate/transpose thing and I am more confident it is correct now, sorry I have overreacted...

jessegrabowski · 2025-03-26T11:52:03Z

Did you check if they actually do fail before you xfailed them? I'm not 100% sure on that

educhesne · 2025-03-26T11:55:50Z

yes I did, verify_grad doesn't accept complex dtype

ricardoV94

Looks great! This is gold

educhesne · 2025-03-29T09:51:45Z

Actually @jessegrabowski even pytensor.gradient.grad doesn't accept complex dtypes.
I found old discussions from the Theano era about it, but nothing else (in particular nothing about choosing the convention for complex gradients, between the jax/holomorphic compatible convention vs the torch/gradient descent compatible one)
Do you know if it is something planned or considered ?

jessegrabowski · 2025-03-29T10:02:31Z

There's been some vague conversations about it, but nothing concrete has been planned. It needs a dev to come along and take up the banner to make it happen

educhesne · 2025-03-29T10:09:45Z

Ok thanks. I could help, but being new to pytensor dev I don't know about being the flag-bearer...

jessegrabowski · 2025-03-29T10:19:29Z

If you're excited about it, or you need it for you work, you should definitely hack around and see if you can put a prototype, and open a draft PR. If you build it, they will come :)

educhesne · 2025-03-29T10:25:41Z

Ok, I'll see... :)
(I don't need it, but that looks interesting...)

jessegrabowski · 2025-03-29T15:19:08Z

Are you a student interested in doing GSoC? We could make a project like that if you are and it's something you're interested in. Or we can brainstorm something else.

educhesne · 2025-03-29T16:18:56Z

Thanks but I am not a student :-D

jessegrabowski · 2025-03-29T16:22:06Z

Ah ok! Sorry I just assumed. We get a lot of activity around this time of year from students looking to participate in GSoC. Anyway still happy to brainstorm on projects/features you'd like to see in pytensor (if you're looking for stuff to do) :)

ricardoV94 · 2025-03-29T17:41:46Z

Thanks but I am not a student :-D

fyi GSOC is no longer restricted to students since last year, it now also accommodates people who are beginning to contribute to OSS in general

Implement gradient for QR decomposition

a187464

jessegrabowski requested changes Mar 19, 2025

View reviewed changes

replace inv by solve_triangular and improve test coverage in qr gradient

ee9aaa2

jessegrabowski approved these changes Mar 19, 2025

View reviewed changes

ricardoV94 requested changes Mar 21, 2025

View reviewed changes

ricardoV94 added gradients linalg labels Mar 21, 2025

Etienne Duchesne added 2 commits March 21, 2025 17:03

qr decomposition gradient: add symbolic shape check

0e47b7d

qr decomposition gradient: extend gradient to more input shapes

9e5e765

* for mode=reduced or mode=r, all input shapes are accepted * for mode=complete, shapes m x n where m <= n are accepted

ricardoV94 reviewed Mar 25, 2025

View reviewed changes

pytensor/tensor/nlinalg.py Outdated Show resolved Hide resolved

qr decompostion gradient: replace swtich by ifelse

ac48c11

jessegrabowski approved these changes Mar 26, 2025

View reviewed changes

qr decomposition gradient: add xfail pytest for complex inputs

a6ae03b

qr decomposition gradient: fix tests

4edc698

ricardoV94 approved these changes Mar 27, 2025

View reviewed changes

ricardoV94 added the enhancement label Mar 27, 2025

ricardoV94 merged commit 2774599 into pymc-devs:main Mar 27, 2025
73 checks passed

educhesne deleted the qr_gradient branch March 28, 2025 07:53

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement gradient for QR decomposition #1303

Implement gradient for QR decomposition #1303

educhesne commented Mar 19, 2025 •

edited by github-actions bot

Loading

jessegrabowski left a comment

jessegrabowski Mar 19, 2025

educhesne commented Mar 19, 2025

jessegrabowski commented Mar 19, 2025 •

edited

Loading

codecov bot commented Mar 19, 2025 •

edited

Loading

educhesne commented Mar 19, 2025

jessegrabowski commented Mar 19, 2025

ricardoV94 left a comment

ricardoV94 Mar 21, 2025

educhesne Mar 21, 2025

ricardoV94 Mar 21, 2025

jessegrabowski left a comment

educhesne commented Mar 26, 2025

jessegrabowski commented Mar 26, 2025

educhesne commented Mar 26, 2025

jessegrabowski commented Mar 26, 2025

educhesne commented Mar 26, 2025 •

edited

Loading

ricardoV94 left a comment

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

ricardoV94 commented Mar 29, 2025

		Rinvt = _H(inv(R))
		A_bar = Q @ ((ptb.tril(R_dRt - _H(R_dRt), k=-1)) @ Rinvt + dR)

Implement gradient for QR decomposition #1303

Implement gradient for QR decomposition #1303

Conversation

educhesne commented Mar 19, 2025 • edited by github-actions bot Loading

Description

Related Issue

Checklist

Type of change

jessegrabowski left a comment

Choose a reason for hiding this comment

jessegrabowski Mar 19, 2025

Choose a reason for hiding this comment

educhesne commented Mar 19, 2025

jessegrabowski commented Mar 19, 2025 • edited Loading

codecov bot commented Mar 19, 2025 • edited Loading

Codecov Report

educhesne commented Mar 19, 2025

jessegrabowski commented Mar 19, 2025

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 Mar 21, 2025

Choose a reason for hiding this comment

educhesne Mar 21, 2025

Choose a reason for hiding this comment

ricardoV94 Mar 21, 2025

Choose a reason for hiding this comment

jessegrabowski left a comment

Choose a reason for hiding this comment

educhesne commented Mar 26, 2025

jessegrabowski commented Mar 26, 2025

educhesne commented Mar 26, 2025

jessegrabowski commented Mar 26, 2025

educhesne commented Mar 26, 2025 • edited Loading

ricardoV94 left a comment

Choose a reason for hiding this comment

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

educhesne commented Mar 29, 2025

jessegrabowski commented Mar 29, 2025

ricardoV94 commented Mar 29, 2025

educhesne commented Mar 19, 2025 •

edited by github-actions bot

Loading

jessegrabowski commented Mar 19, 2025 •

edited

Loading

codecov bot commented Mar 19, 2025 •

edited

Loading

educhesne commented Mar 26, 2025 •

edited

Loading