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,