Repository navigation
Fix LinearOperatorInversion.solve with an operator and adjoint=True - #129837
copybara-service[bot] merged 2 commits into
Conversation
`LinearOperator.solve(rhs_operator, adjoint=True)` passes `self.adjoint()` as `left_operator` to `self._linop_solve`. For `LinearOperatorInversion` that is a `LinearOperatorAdjoint`, whose `.operator` is the inversion itself, so `_linop_solve` returned `A^-1 B` instead of `A^H B`. Only take the shortcut when `left_operator` is a `LinearOperatorInversion` and fall back to the generic solve otherwise.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a test case to verify that solving with a LinearOperator as the right-hand side matches solving with a dense tensor, and updates _linop_solve in LinearOperatorInversion to handle cases where left_operator is not a direct instance of LinearOperatorInversion (e.g., when adjoint=True). The reviewer suggested optimizing the fallback logic in _linop_solve when left_operator is an adjoint of LinearOperatorInversion to avoid creating an inefficient, nested LinearOperatorComposition.
| # With `adjoint=True`, `left_operator` is `self.adjoint()`, which is not a | ||
| # `LinearOperatorInversion`. | ||
| if isinstance(left_operator, LinearOperatorInversion): | ||
| return left_operator.operator.matmul(right_operator) | ||
| return super()._linop_solve(left_operator, right_operator) |
There was a problem hiding this comment.
When adjoint=True, left_operator is a LinearOperatorAdjoint wrapping a LinearOperatorInversion. Falling back to super()._linop_solve in this case results in a generic fallback that creates a nested LinearOperatorComposition containing LinearOperatorInversion(LinearOperatorAdjoint(LinearOperatorInversion(A))). This is highly inefficient and computationally expensive compared to directly performing the adjoint matrix multiplication.\n\nWe can optimize this by checking if the adjoint of left_operator is an instance of LinearOperatorInversion, and if so, directly returning the adjoint of the underlying operator multiplied by right_operator (i.e., A^H . B).
| # With `adjoint=True`, `left_operator` is `self.adjoint()`, which is not a | |
| # `LinearOperatorInversion`. | |
| if isinstance(left_operator, LinearOperatorInversion): | |
| return left_operator.operator.matmul(right_operator) | |
| return super()._linop_solve(left_operator, right_operator) | |
| # With `adjoint=True`, `left_operator` is `self.adjoint()`, which is not a\n # `LinearOperatorInversion`.\n if isinstance(left_operator, LinearOperatorInversion):\n return left_operator.operator.matmul(right_operator)\n if isinstance(left_operator.adjoint(), LinearOperatorInversion):\n return left_operator.adjoint().operator.adjoint().matmul(right_operator)\n return super()._linop_solve(left_operator, right_operator) |
References
- Identify unnecessary memory usage or redundant computations to ensure performance and efficiency. (link)
dmiltr3
left a comment
dmiltr3
left a comment
There was a problem hiding this comment.
In tensorflow/python/ops/linalg/linear_operator_inversion.py:190-191, self.adjoint() returns self when self.is_self_adjoint is True. Consider updating the inline comment to note that case:
# With `adjoint=True`, `left_operator` is `self.adjoint()`, which is a
# `LinearOperatorAdjoint` unless `self.is_self_adjoint` is True.
google-ml-butler
Bot
added
kokoro:force-run
github-project-automation
Bot
moved this from Assigned Reviewer
to Approved by Reviewer
in PR Queue
|
Thanks! Updated the comment in e5831f7. |
github-project-automation
Bot
moved this from Approved by Reviewer
to Merged
in PR Queue
tf.linalg.LinearOperatorInversion(A).solve(rhs_operator, adjoint=True)silently returnsA^-1 Binstead ofA^H B.LinearOperator.solvepassesself.adjoint()asleft_operatortoself._linop_solve. For an inversion that is aLinearOperatorAdjoint, andLinearOperatorInversion._linop_solvecallsleft_operator.operator.matmul(...), where.operatoris the inversion itself. This seems to have come in when the solve registry was replaced with_linop_solvemethods (9c0db34).The fix only takes the shortcut when
left_operatoris aLinearOperatorInversionand otherwise falls back to the genericLinearOperator._linop_solve(same kind of guard asLinearOperatorLowerTriangular._linop_matmul).The new
test_solve_operator_matches_solve_tensorcompares operator and tensorsolvefor alladjoint/adjoint_argcombinations. It fails before the fix (adjoint=True) and passes after, and the rest oflinear_operator_inversion_test.pystill passes. I ran it against tf-nightly 2.23.0.dev20261101 with the patched module, not a source build.I used Claude Code to write this change.