Add rewrite to merge multiple SVD Ops with different settings

.gitignore

@@ -56,3 +56,5 @@ pytensor-venv/
 testing-report.html
 coverage.xml
 .coverage.*
+pics
+*.ipynb

pytensor/tensor/rewriting/linalg.py

@@ -4,13 +4,17 @@
 
 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,
@@ -377,3 +381,26 @@
                 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
+    compute_uv = False
+
+    for cl, _ in fgraph.clients[x]:
+        if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+            if (not compute_uv) and cl.op.core_op.compute_uv:
+                compute_uv = True
+                break
+
+    if compute_uv and not node.op.compute_uv:
+        full_matrices = node.op.full_matrices
+        return [SVD(full_matrices=full_matrices, compute_uv=compute_uv)]