pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Implement generalized rewrite to replace ops that compute partial outputs by ops with all the outputs

Open ricardoV94 opened this issue 9 months ago • 0 comments

Description

Several linalg methods like SVD and QR allow to only request a subset of the outputs.

If we have multiple of these with the same input and different output subsets we should use the one that computes the whole subset once. Also, in the gradient we usually need all the outputs anyway.

It's perhaps easier if we ignore the user request and the user-facing function always uses the Op with all the outputs (not returning the ones not requested by the user). In the specialize rewrite we can specialize into the smaller subset Op. By this time gradients should already have been requested.

We can make the gradient raise if called on an Op which is not computing all the outputs and thus simplify the grad method as well.

ricardoV94 avatar Mar 21 '25 11:03 ricardoV94