Add rewrite to remove transposed kwargs

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -49,6 +49,7 @@
     Cholesky,
     Solve,
     SolveBase,
+    SolveTriangular,
     _bilinear_solve_discrete_lyapunov,
     block_diag,
     cholesky,
@@ -957,7 +958,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 +1015,45 @@ 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_transposed_argument_to_symbolic(fgraph, node):
+    """
+    Replace solve(A, b, transposed=True) with solve(A.T, b, transposed=False). In testing this was slightly faster,
+    and results is more readable graphs, because the transpose Op is explicitly included.
+    """
+    solve_op = node.op.core_op
+    if not isinstance(solve_op, Solve | SolveTriangular):
+        return None
+
+    solve_kwargs = solve_op._props_dict()
+    A, b = node.inputs
+
+    if isinstance(solve_op, Solve):
+        if not solve_kwargs["transposed"]:
+            return None
+
+        solve_kwargs["transposed"] = False
+        solve_kwargs["lower"] = not solve_kwargs["lower"]
+        return [Blockwise(Solve(**solve_kwargs))(A.T, b)]
+
+    if isinstance(solve_op, SolveTriangular):
+        if solve_kwargs["trans"] in [False, 0]:
+            return None
+
+        elif solve_kwargs["trans"] in [True, 1, "T"]:
+            solve_kwargs["trans"] = 0
+            solve_kwargs["lower"] = not solve_kwargs["lower"]
+            return [Blockwise(SolveTriangular(**solve_kwargs))(A.T, b)]
+
+        elif solve_kwargs["trans"] in [2, "C"]:
+            solve_kwargs["trans"] = 0
+            solve_kwargs["lower"] = not solve_kwargs["lower"]
+            return [Blockwise(SolveTriangular(**solve_kwargs))(A.conj().T, b)]
+
+        else:
+            return None
+
+    return None

pytensor/tensor/slinalg.py

@@ -284,6 +284,13 @@ def L_op(self, inputs, outputs, output_gradients):
 
         Symbolic expression for updates taken from [#]_.
 
+        .. math::
+
+            \begin{align}
+            \bar{b} &= A^{-T} \bar{c} \\
+            \bar{A} &= -A^{-T} \bar{c} c^T &= \bar{b} \c^T \\
+            \end{align}
+
         References
         ----------
         .. [#] M. B. Giles, "An extended collection of matrix derivative results
@@ -294,31 +301,24 @@ def L_op(self, inputs, outputs, output_gradients):
         A, b = inputs
 
         c = outputs[0]
-        # C is a scalar representing the entire graph
-        # `output_gradients` is (dC/dc,)
-        # We need to return (dC/d[inv(A)], dC/db)
         c_bar = output_gradients[0]
 
         props_dict = self._props_dict()
-
+        # Note that we do *not* reverse the lower argument here. This is because LAPACK handles this internally when
+        # transpose=True. If we try to "help", it will end up solving with the wrong triangle.
         if isinstance(self, SolveTriangular):
             # SolveTriangular has a special trans argument we have to handle
-            transposed = props_dict.pop("trans") in [1, "T"]
+            transposed = props_dict.pop("trans") in [1, "T", 2, "C"]
             props_dict["trans"] = not transposed
         else:
             transposed = props_dict.pop("transposed")
             props_dict["transposed"] = not transposed
 
-        # TODO: We were flipping lower before, but it doesn't appear we need to -- all tests pass without taking it into
-        #  account.
-        # props_dict['lower'] = not self.lower
-
         solve_op = type(self)(**props_dict)
-
         b_bar = solve_op(A, c_bar)
+
         # force outer product if vector second input
         A_bar = -ptm.outer(b_bar, c) if c.ndim == 1 else -b_bar.dot(c.T)
-
         if transposed:
             A_bar = A_bar.T
 

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,35 @@ 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)
+
+
+@pytest.mark.parametrize(
+    "solve_op, kwargs",
+    [
+        (pt.linalg.solve, {"transposed": True, "lower": False, "assume_a": "sym"}),
+        (pt.linalg.solve_triangular, {"trans": 1, "lower": False}),
+    ],
+)
+def test_rewrite_away_transposed_argument(solve_op, kwargs):
+    A = pt.tensor("A", shape=(5, 5))
+    b = pt.tensor("b", shape=(5,))
+    x = solve_op(A, b, **kwargs)
+
+    fg = FunctionGraph([A, b], [x])
+    assert not any(isinstance(node.op, DimShuffle) for node in fg.toposort())
+
+    f = function([A, b], x, mode="FAST_RUN")
+    assert any(isinstance(node.op, DimShuffle) for node in f.maker.fgraph.toposort())
+
+    A_val = np.triu(np.random.normal(size=(5, 5)))
+    b_val = np.random.normal(size=(5,))
+
+    new_kwargs = kwargs.copy()
+    if "transposed" in kwargs:
+        new_kwargs["transposed"] = not kwargs["transposed"]
+    if "trans" in kwargs:
+        new_kwargs["trans"] = int(not bool(kwargs["trans"]))
+    new_kwargs["lower"] = not kwargs["lower"]
+
+    g = function([A, b], solve_op(A.T, b, **new_kwargs), mode="FAST_COMPILE")
+    np.testing.assert_allclose(f(A_val, b_val), g(A_val, b_val))

tests/tensor/test_slinalg.py

@@ -332,13 +332,18 @@ def test_solve_gradient(
 
         def A_func(x):
             if assume_a == "pos":
-                return x @ x.T
+                x = x @ x.T
             elif assume_a == "sym":
-                return (x + x.T) / 2
-            else:
-                return x
+                x = (x + x.T) / 2
+            return x
 
-        solve_op = functools.partial(solve, assume_a=assume_a, b_ndim=len(b_size))
+        solve_op = functools.partial(
+            solve,
+            assume_a=assume_a,
+            lower=lower,
+            transposed=transposed,
+            b_ndim=len(b_size),
+        )
 
         # To correctly check the gradients, we need to include a transformation from the space of unconstrained matrices
         # (A) to a valid input matrix for the given solver. This is done by the A_func function. If this isn't included,