Skip to content

Collapse quadratic forms to a single triangular solve - #2373

Open
jessegrabowski wants to merge 2 commits into
pymc-devs:mainfrom
jessegrabowski:quadratic-form-single-solve
Open

Collapse quadratic forms to a single triangular solve#2373
jessegrabowski wants to merge 2 commits into
pymc-devs:mainfrom
jessegrabowski:quadratic-form-single-solve

Conversation

@jessegrabowski

@jessegrabowski jessegrabowski commented Aug 23, 2026

Copy link
Copy Markdown
Member

Computing b.T @ K^-1 @ b applies K^-1 with a forward and a back substitution, and then the surrounding dot contracts the result straight back to a scalar. Writing K = C C.T, the same quantity is ||C^-1 b||^2, so one substitution does. The rewrite catches cho_solve and also a hand-written pair of transposed triangular solves, whose factor doesn't have to come from a Cholesky; the square-root Kalman filter gets its factor out of a QR.

The first commit lifts the b_ndim=2 gate on psd_solve_to_chol_solve so a positive-definite solve with a vector right-hand side reaches the same chain instead of staying one opaque posv call.

The two benchmarks differ in whether the triangular factor is already available. In given_factor it is, and dropping one substitution roughly halves the time. In with_cholesky the factorization is redone on every call, so its O(n^3) cost dominates a single quadratic form and the saving shrinks to 6-8%, which is within one standard deviation except at n=3000.

------------------------------------ benchmark 'n=1000': 4 tests -------------------------------------
Name (time in us)                                         Median              StdDev            Rounds
------------------------------------------------------------------------------------------------------
test_quadratic_form_given_factor[1000-baseline]         281.9160 (2.08)      11.5286 (1.49)       3444
test_quadratic_form_given_factor[1000-rewritten]        135.7915 (1.0)        7.7555 (1.0)        7107
test_quadratic_form_with_cholesky[1000-baseline]      2,174.2499 (16.01)    117.2891 (15.12)       463
test_quadratic_form_with_cholesky[1000-rewritten]     2,044.7075 (15.06)    128.6777 (16.59)       509
------------------------------------------------------------------------------------------------------

----------------------------------- benchmark 'n=200': 4 tests ----------------------------------
Name (time in us)                                     Median             StdDev            Rounds
-------------------------------------------------------------------------------------------------
test_quadratic_form_given_factor[200-baseline]       15.3743 (1.85)      2.0301 (1.74)      46155
test_quadratic_form_given_factor[200-rewritten]       8.2916 (1.0)       1.1643 (1.0)       17532
test_quadratic_form_with_cholesky[200-baseline]      60.7083 (7.32)     18.1682 (15.60)      9088
test_quadratic_form_with_cholesky[200-rewritten]     56.1252 (6.77)      6.3358 (5.44)      15545
-------------------------------------------------------------------------------------------------

---------------------------------- benchmark 'n=3000': 4 tests ----------------------------------
Name (time in ms)                                      Median            StdDev            Rounds
-------------------------------------------------------------------------------------------------
test_quadratic_form_given_factor[3000-baseline]        3.0421 (1.90)     0.2504 (1.47)        329
test_quadratic_form_given_factor[3000-rewritten]       1.6029 (1.0)      0.1705 (1.0)         571
test_quadratic_form_with_cholesky[3000-baseline]      26.6957 (16.65)    0.5196 (3.05)         50
test_quadratic_form_with_cholesky[3000-rewritten]     24.6633 (15.39)    0.8209 (4.81)         50
-------------------------------------------------------------------------------------------------

@jessegrabowski jessegrabowski added enhancement New feature or request linalg Linear algebra graph rewriting labels Aug 23, 2026
@jessegrabowski
jessegrabowski force-pushed the quadratic-form-single-solve branch from 5d33e7c to 477ab59 Compare August 26, 2026 01:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request graph rewriting linalg Linear algebra

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant