Refactor and update QR Op

[jessegrabowski]

Description

This PR updates the QRFull Op, adding static shape checking, infer_shape, and destroy_map. It also optimizes the perform method for the C backend, and tries to improve the gradient graph by checking static shapes (to avoid an ifelse).

I renamed it to QR, because I don't know what was Full about the old one. I also moved it from the numpy implementation to scipy, which gives us all the usual benefits (inplace, etc). I also went ahead and unpacked the scipy wrapper and used the LAPACK functions directly. This will give us better error handling (that is to say, none -- it should eventually return a matrix of NaN on failure) and some performance boost by caching workspace requirements.

Still a WIP, because it breaks everything by moving QR from nlinalg to slinalg. I thought about using this as an opportunity to finally eliminate this distinction and go to a more logical organization (linalg/decomposition/qr.py), but then decided against it for now. Needs discussion.

Related Issue

• Closes Add infer_shape method to QRFull #1511
• Related to #

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

---

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

[ricardoV94]

It's better not to give names to intermediate variables, this would shadow a user Q, if they wanted to do eval({"Q": something}) or find a variable by name

[ricardoV94]

You shouldn't mutate an Op, it can be reused multiple times. If you want, mutate the tag of an apply node, since that's unique in a graph, whereas an Op is not.

[jessegrabowski]

yeah I know this is dumb and I was hoping to push a change before you arrived with this comment. Mission failed.

[ricardoV94]

Put this in make_node, the perform shouldn't really get anything the make_node didn't allow

[ricardoV94]

Instead of duplicating code logic can you add a eager_ifelse? I did something like that in this WIP PR: 0dc2a1e

def _eager_switch(
    cond: TensorVariable | bool, a: TensorVariable, b: TensorVariable
) -> TensorVariable:
    # Do not create a switch if cond is True/False
    # We need this because uint types cannot be negative and creating the lazy switch could upcast everything to float64
    # It also simplifies immediately the graph that's returned
    if isinstance(cond, bool):
        return a if cond else b
    return cast(TensorVariable, switch(cond, a, b))

This way the code is the same whether you can do eagerly or not? Or do you actually manage to return something different in the branches if you know in advance?

[jessegrabowski]

how would I use the symbolic shapes as cond in this case? Is there a point during rewrites where shapes are evaluated if possible?

That is to say, I did it this way because I thought I needed to check the static shape and dynamic shape separately

[ricardoV94]

I see is it a problem if the ifelse only disappears during constant fold in compilation?

If yes I can bother with thinking of a solution

[jessegrabowski]

no that's fine. Is that what would happen with _eager_switch? I don't see how if so

[jessegrabowski]

to be clear though there's no logic duplication here, only the extra returns inside each if block. The QR gradient always has two cases, depending on the shape of the input matrix

[ricardoV94]

There's some complexity because you're trying to use static shapes and prune branches of the ifelse eagerly. PyTensor will figure out easily enough M > N in an ifelse if it has static shapes and prune the useless ifelse branch.

If there's a reason why you can't wait for pytensor to do this (in the subtensor PR it just lead to a giant graph complexity) we can try to create an eager_ifelse that tries to see if M > N and only return one branch if that's the case, so that the L_Op code looks pretty much the same regardless of whether it can be removed or not.