Added rewrite for matrix inv(inv(x)) -> x

[tanish1729]

Description

Adds rewrite for inv(inv(x)) -> x

Related Issue

• Closes Rewrite for consecutive matrix inverses #887

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[jessegrabowski]

Great start! Biggest missing piece is the solve-based inverses. Let's use this opportunity to write a helper that detects those, because all inverse rewrites will need it.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.50%. Comparing base (ad27dc7) to head (4bf3c2d).
Report is 91 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #893   +/-   ##
=======================================
  Coverage   81.49%   81.50%           
=======================================
  Files         176      176           
  Lines       46925    46938   +13     
  Branches    11428    11435    +7     
=======================================
+ Hits        38242    38255   +13     
- Misses       6498     6500    +2     
+ Partials     2185     2183    -2     

Files with missing lines	Coverage Δ
pytensor/tensor/rewriting/linalg.py	91.28% <100.00%> (+0.49%)	⬆️

... and 2 files with indirect coverage changes

[jessegrabowski]

lgtm

[tanish1729]

@ricardoV94 lmk if there's anything else to be done here. otherwise, it should be good to merge

[ricardoV94]

If we predefine the two possible pinv as blockwise (like we do for matrix_inverse) the rewrite can track more specifically and avoids being called for any Blockwise it sees:

Suggested change

	@node_rewriter([Blockwise])
	def rewrite_inv_inv(fgraph, node):
	"""
	This rewrite takes advantage of the fact that if there are two consecutive inverse operations (inv(inv(input))), we get back our original input without having to compute inverse once.
	Here, we check for direct inverse operations (inv/pinv) and allows for any combination of these "inverse" nodes to be simply rewritten.
	Parameters
	----------
	fgraph: FunctionGraph
	Function graph being optimized
	node: Apply
	Node of the function graph to be optimized
	Returns
	-------
	list of Variable, optional
	List of optimized variables, or None if no optimization was performed
	"""
	valid_inverses = (MatrixInverse, MatrixPinv)
	# Check if its a valid inverse operation (either inv/pinv)
	# In case the outer operation is an inverse, it directly goes to the next step of finding inner operation
	# If the outer operation is not a valid inverse, we do not apply this rewrite
	if not isinstance(node.op.core_op, valid_inverses):
	return None
	@node_rewriter([matrix_inverse, matrix_pinv_hermitian, matrix_pinv_non_hermitian])
	def rewrite_inv_inv(fgraph, node):
	"""
	This rewrite takes advantage of the fact that if there are two consecutive inverse operations (inv(inv(input))), we get back our original input without having to compute inverse once.
	Here, we check for direct inverse operations (inv/pinv) and allows for any combination of these "inverse" nodes to be simply rewritten.
	Parameters
	----------
	fgraph: FunctionGraph
	Function graph being optimized
	node: Apply
	Node of the function graph to be optimized
	Returns
	-------
	list of Variable, optional
	List of optimized variables, or None if no optimization was performed
	"""

Need to predefine those matrix_pinv*.. The helper "pinv" should return the predefined Ops instead of creating new ones to avoid Op duplication

[tanish1729]

i dont quite understand the last line u wrote about

Need to predefine those matrix_pinv*.. The helper "pinv" should return the predefined Ops instead of creating new ones to avoid Op duplication

also, as a general rule for rewrites, is it better that they should be tracking more specific Ops instead of just Blockwise

[ricardoV94]

also, as a general rule for rewrites, is it better that they should be tracking more specific Ops instead of just Blockwise

Is that a question? The answer is yes. It avoids useless calls to the rewrite function when the Op is not in the graph

[tanish1729]

but for the check of the ops inside, i will have to use the method that i am already doing

[ricardoV94]

Regarding pinv, the helper function is here:

pytensor/pytensor/tensor/nlinalg.py

Line 68 in ad27dc7

def pinv(x, hermitian=False):

The idea is instead:

pinv_hermitian = Blockwise(Pinv(hermitian=True))
pinv_non_hermitian = Blockwise(Pinv(hermitian=False))

def pinv(x, hermitian=False):
  ...
  return pinv_hermitian if hermitian else pinv_non_hermitian

[ricardoV94]

But actually let's not to that now, have to think. This makes initialization a bit slower because we have to create more instances...

[tanish1729]

oh alright. but this is essentially the same thing as first checking the op as blockwise and then the core op as one of the valid inverses. whats the difference in both these ideas

[ricardoV94]

The difference is the rewrite is not even considered if the node is not a Blockwise(Pinv). This way the rewrite is attempted on every single Blockwise in the graph.

It's just an optimization, not a question of correctness