Implement generalized rewrite to replace ops that compute partial outputs by ops with all the outputs #1312
Description
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.