pymc-devs / pymc-devs/pytensor

Matrix determinant lemma example shows several optimizations misses

Open
#1,040 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

graph rewriting linalg
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.