Add MultiDot op and rewrites for optimal contraction

pytensor/tensor/linalg/products.py

@@ -1,3 +1,5 @@
+import warnings
+
 import scipy.linalg as scipy_linalg
 
 from pytensor import tensor as pt
@@ -126,9 +128,15 @@ def matrix_dot(*args):
     :math:`A_0 \cdot A_1 \cdot A_2 \cdot .. \cdot A_N`.
 
     """
+    warnings.warn(
+        "matrix_dot is deprecated and will be removed in future version.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
     rval = args[0]
     for a in args[1:]:
-        rval = ptm.dot(rval, a)
+        rval = ptm.matmul(rval, a)
     return rval
 
 

pytensor/tensor/rewriting/linalg/__init__.py

@@ -6,6 +6,7 @@
 import pytensor.tensor.rewriting.linalg.decomposition
 import pytensor.tensor.rewriting.linalg.inverse
 import pytensor.tensor.rewriting.linalg.products
+import pytensor.tensor.rewriting.linalg.reassociate_matmul
 import pytensor.tensor.rewriting.linalg.solvers
 import pytensor.tensor.rewriting.linalg.summary
 import pytensor.tensor.rewriting.linalg.utils

pytensor/tensor/rewriting/linalg/products.py

@@ -1,7 +1,4 @@
-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 ExtractDiag, concatenate, diag
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.linalg.constructors import BlockDiagonal