Add rewrite to merge multiple SVD Ops with different settings

.gitignore

@@ -55,4 +55,4 @@ pytensor-venv/
 .vscode/
 testing-report.html
 coverage.xml
-.coverage.*
+.coverage.*

pytensor/tensor/rewriting/linalg.py

@@ -4,20 +4,25 @@
 
 from pytensor import Variable
 from pytensor.graph import Apply, FunctionGraph
-from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
+from pytensor.graph.rewriting.basic import (
+    copy_stack_trace,
+    node_rewriter,
+)
 from pytensor.tensor.basic import TensorVariable, diagonal
 from pytensor.tensor.blas import Dot22
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.math import Dot, Prod, _matrix_matrix_matmul, log, prod
 from pytensor.tensor.nlinalg import (
+    SVD,
     KroneckerProduct,
     MatrixInverse,
     MatrixPinv,
     det,
     inv,
     kron,
     pinv,
+    svd,
 )
 from pytensor.tensor.rewriting.basic import (
     register_canonicalize,
@@ -377,3 +382,47 @@
                 return [block_diag(*inner_matrices)]
             else:
                 raise NotImplementedError  # pragma: no cover
+
+
+@register_canonicalize
+@register_stabilize
+@register_specialize
+@node_rewriter([SVD])
+def local_svd_uv_simplify(fgraph, node):
+    """If we have more than one `SVD` `Op`s and at least one has keyword argument
+    `compute_uv=True`, then we can change `compute_uv = False` to `True` everywhere
+    and allow `pytensor` to re-use the decomposition outputs instead of recomputing.
+    """
+    (x,) = node.inputs
+
+    if node.compute_uv:
+        # compute_uv=True returns [u, s, v].
+        # if at least u or v is used, no need to rewrite this node.
+        if (
+            fgraph.clients[node.outputs[0]] is not None
+            or fgraph.clients[node.outputs[2]] is not None
+        ):
+            return
+
+        # Else, has to replace the s of this node with s of an SVD Op that compute_uv=False.
+        # First, iterate to see if there is an SVD Op that can be reused.
+        for cl, _ in fgraph.clients[x]:
+            if cl == "output":
+                continue
+            if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+                if not cl.op.core_op.compute_uv:
+                    return {fgraph.clients[node.outputs[1]]: cl.outputs[0]}
+
+        # If no SVD reusable, return a new one.
+        return [svd(x, full_matrices=node.full_matrices, compute_uv=False)]
+
+    else:
+        # compute_uv=False returns [s].
+        # We want rewrite if there is another one with compute_uv=True.
+        # For this case, just reuse the `s` from the one with compute_uv=True.
+        for cl, _ in fgraph.clients[x]:
+            if cl == "output":
+                continue
+            if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+                if cl.op.core_op.compute_uv:
+                    return [cl.outputs[1]]