psd_solve_to_chol_solve only tracks b_ndim=2, so solve(A, b, assume_a="pos") with a vector b factors A twice
Nobody has claimed this yet.
Assessment
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Newbie friendliness
- 78/100
Research direction
Start by locating the psd_solve_to_chol_solve stabilize rewrite and its OpPattern(Solve, b_ndim=2) registration. Extend its coverage to vector right-hand sides while preserving the existing matrix case, then verify that the example with b_ndim=1 shares one Cholesky factor instead of factoring A twice.
Written by the indexing model from the issue text.
Description
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
- Dominant language
- Python
- Stars
- 644
- Forks
- 209
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from pymc-devs/pytensor
-
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
-
bug indexing mlx
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 84/100
-
bug torch
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
All issues in pymc-devs/pytensor
Similar issues
-
Difficulty 2/5 1-3 hours Newbie friendliness 88/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
-
enhancement
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 74/100