Skip to content

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

Open
@ricardoV94

Description

@ricardoV94

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions