use mv in ComposedLinearOperator.as_matrix to avoid unecessary materialisation#196
Open
jpbrodrick89 wants to merge 3 commits intopatrick-kidger:mainfrom
Open
use mv in ComposedLinearOperator.as_matrix to avoid unecessary materialisation#196jpbrodrick89 wants to merge 3 commits intopatrick-kidger:mainfrom
jpbrodrick89 wants to merge 3 commits intopatrick-kidger:mainfrom
Conversation
- Add identity shortcuts: Identity @ X returns X, X @ Identity returns X - Use vmap over operator1.mv instead of matmul for as_matrix, enabling efficient composition when operator1 has O(N) mv (e.g., Diagonal) https://claude.ai/code/session_0143xm3Fot5bh7Zy3GfkP6bD
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I ran into a case where I wanted to do
DiagonalLinearOperator @ TridiagonalLinearOperatorto avoid having to write three lines of code but realised that the current implementation would make this O(n^3) matrix-by-matrix multiplication instead of O(n). This change takes advantage of efficientmvin the outer operator which makes the operation ~O(n^2). We can work towards making it O(n) as part of the tridiagonal PR. (For example, the diagonal PR handle this by using try_sparse_materialise).Note there are loads of edge cases here, especially if
op2has an efficient mv butop1doesn't, in such cases it may be more efficient to call(operator2.T.mv(op1.T.as_matrix()).Tand we could achieve this by creating a tuple heirarchy and reading off the index. However, this simpler implementation should match user's expectation of what is going on under the hood and they probably wouldn't be surprised to find thatMatrixLinearOperator @ DiagonalLinearOperatoris inefficient and accept that they should manually handle such cases themselves.I would also be willing to remove the special handling I added for
IdentityLinearOperatorfor a similar reason: it is very unlikely that as user is going to decide compose an IdentityLinearOperator with another.