pymc-devs / pymc-devs/pytensor
Matrix determinant lemma example shows several optimizations misses
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
https://en.wikipedia.org/wiki/Matrix_determinant_lemma
Ideally, in the example below we would rewrite from the naive out to opt_out, when we can be confident that computing the determinant of A plus solving v.T A^-1 u should be cheaper than computing the determinant of B.
In practice, this may require doing some tentative rewrites and potentially evaluating flops (#875 would be nice), and backtracking if we don't find anything useful for the det of A.
The example below is setup so this is the case, by making A a diagonal matrix.
import numpy as np
import pytensor
import pytensor.tensor as pt
a = pt.vector("a")
A = pt.diag(a)
u = pt.vector("u")
v = pt.vector("v")
B = A + pt.outer(u, v)
out = pt.linalg.det(B)
fn = pytensor.function([a, u, v], out)
alt_out = (1 + v @ pt.linalg.solve(A, u)) * pt.linalg.det(A)
alt_fn = pytensor.function([a, u, v], alt_out)
opt_out = (1 + v @ ((1 / a) * u)) * pt.prod(a)
opt_fn = pytensor.function([a, u, v], opt_out)
rng = np.random.default_rng()
a_test, u_test, v_test = rng.normal(size=(3, 40))
np.testing.assert_allclose(
fn(a_test, u_test, v_test),
alt_fn(a_test, u_test, v_test),
)
np.testing.assert_allclose(
fn(a_test, u_test, v_test),
opt_fn(a_test, u_test, v_test),
)
%timeit fn(a_test, u_test, v_test) # 49.2 μs ± 1.45 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit alt_fn(a_test, u_test, v_test) # 110 μs ± 10.7 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit opt_fn(a_test, u_test, v_test) # 17.9 μs ± 1.05 μs per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
The intermediate form alt_out highlights some extra missing rewrites, such as the diagonal solve and vector @ diagonal matrix products.
The matrix determinant lemma equivalence can be further beneficial if the graph also requires a solve/inverse of B, as some of the new expressions can be reused following: https://en.wikipedia.org/wiki/Sherman%E2%80%93Morrison_formula
Another cases besides diagonal A, would be triangular or circulant matrices. There's a nice blog on the latter.
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.
Research direction
No source file or test is named. Start by reproducing the Python determinant example and inspect the existing determinant, solve, diagonal, and graph-rewrite entry points; done means equivalent expressions are optimized without unsafe rewrites, with coverage for the demonstrated case and the stated matrix forms.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 28/100