psd_solve_to_chol_solve only tracks b_ndim=2, so solve(A, b, assume_a="pos") with a vector b factors A twice

[jessegrabowski]

pt.linalg.solve(A, b, assume_a="pos", b_ndim=1) next to pt.linalg.cholesky(A) compiles to a Solve and a Cholesky of the same matrix. psd_solve_to_chol_solve (the stabilize rewrite that turns a positive-definite solve into cholesky plus triangular solves, which is what lets MergeOptimizer share the factor) tracks OpPattern(Solve, b_ndim=2) only, so a vector right-hand side never reaches it. solve_triangular accepts b_ndim=1, so the rewrite could track both.

import pytensor
import pytensor.tensor as pt

A = pt.matrix("A", shape=(5, 5))
b = pt.vector("b", shape=(5,))

logdet = 2 * pt.log(pt.diagonal(pt.linalg.cholesky(A))).sum()
quad = b @ pt.linalg.solve(A, b, assume_a="pos", b_ndim=1)

pytensor.dprint(pytensor.function([A, b], logdet + quad))
# Solve{assume_a='pos', b_ndim=1}(A, b) and Cholesky(A): A is factored twice
# workaround: solve(A, b[:, None], assume_a="pos", b_ndim=2)[:, 0] -> one Cholesky, CholeskySolve