Skip to content

use mv in ComposedLinearOperator.as_matrix to avoid unecessary materialisation#196

Open
jpbrodrick89 wants to merge 3 commits intopatrick-kidger:mainfrom
jpbrodrick89:jpb/composed-materialise
Open

use mv in ComposedLinearOperator.as_matrix to avoid unecessary materialisation#196
jpbrodrick89 wants to merge 3 commits intopatrick-kidger:mainfrom
jpbrodrick89:jpb/composed-materialise

Conversation

@jpbrodrick89
Copy link
Contributor

I ran into a case where I wanted to do DiagonalLinearOperator @ TridiagonalLinearOperator to 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 efficient mv in 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 op2 has an efficient mv but op1 doesn't, in such cases it may be more efficient to call (operator2.T.mv(op1.T.as_matrix()).T and 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 that MatrixLinearOperator @ DiagonalLinearOperator is inefficient and accept that they should manually handle such cases themselves.

I would also be willing to remove the special handling I added for IdentityLinearOperator for a similar reason: it is very unlikely that as user is going to decide compose an IdentityLinearOperator with another.

jpbrodrick89 and others added 3 commits February 2, 2026 12:13
- 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant