Add rewrite to move transpose operations to BLAS

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -957,7 +957,8 @@ def jax_bilinaer_lyapunov_to_direct(fgraph: FunctionGraph, node: Apply):
 @node_rewriter([det])
 def slogdet_specialization(fgraph, node):
     """
-    This rewrite targets specific operations related to slogdet i.e sign(det), log(det) and log(abs(det)) and rewrites them using the SLogDet operation.
+    This rewrite targets specific operations related to slogdet i.e sign(det), log(det) and log(abs(det)) and rewrites
+    them using the SLogDet operation.
 
     Parameters
     ----------
@@ -1013,3 +1014,30 @@ def slogdet_specialization(fgraph, node):
             k: slogdet_specialization_map[v] for k, v in dummy_replacements.items()
         }
         return replacements
+
+
+@register_specialize
+@node_rewriter([Blockwise])
+def rewrite_A_transposed_solve_to_transposed_solver(fgraph, node):
+    """
+    Replace solve(A.T, b) with solve(A, b, transposed=True).
+    """
+    solve_op = node.op.core_op
+    if not isinstance(solve_op, Solve):
+        return None
+
+    A, b = node.inputs
+
+    if not is_matrix_transpose(A):
+        return None
+
+    # If we've gotten here, A is actually A.T
+    A = A.owner.inputs[0]
+    solve_kwargs = solve_op._props_dict()
+
+    # The reason for using `not transposed` here is that we could have had solve(A.T, b, transposed=True), in which
+    # case we can just do solve(A, b, transposed=False) and the rewrite is still valid. In the "base case" we had
+    # solve(A.T, b, transposed=False) and we're going to solve(A, b, tranposed=True).
+    solve_kwargs["transposed"] = not solve_kwargs["transposed"]
+
+    return [Blockwise(Solve(**solve_kwargs))(A, b)]

tests/tensor/rewriting/test_linalg.py

@@ -10,6 +10,7 @@
 from pytensor import tensor as pt
 from pytensor.compile import get_default_mode
 from pytensor.configdefaults import config
+from pytensor.graph import FunctionGraph
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor import swapaxes
 from pytensor.tensor.blockwise import Blockwise
@@ -993,3 +994,23 @@ def test_slogdet_specialization():
     f = function([x], [exp_det_x, sign_det_x], mode="FAST_RUN")
     nodes = f.maker.fgraph.apply_nodes
     assert not any(isinstance(node.op, SLogDet) for node in nodes)
+
+
+def test_rewrite_A_transposed_solve_to_transposed_solver():
+    A = matrix("A")
+    b = vector("b")
+    x = pt.linalg.solve(A.T, b)
+
+    fg = FunctionGraph([A, b], [x])
+    assert any(isinstance(node.op, DimShuffle) for node in fg.toposort())
+
+    f = function([A, b], x, mode="FAST_RUN")
+    assert not any(
+        isinstance(node.op, DimShuffle) for node in f.maker.fgraph.toposort()
+    )
+
+    A_val = np.random.normal(size=(10, 10))
+    b_val = np.random.normal(size=(10,))
+
+    g = function([A, b], pt.linalg.solve(A.T, b), mode="FAST_COMPILE")
+    np.testing.assert_allclose(f(A_val, b_val), g(A_val, b_val))