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

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.

given_factor times the substitutions with the factor already available; with_cholesky factorizes from the covariance matrix on every call, so one quadratic form has to amortize the whole O(n^3) and only 6-8% is left, inside a 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
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