pymc-devs / pymc-devs/pytensor

Add rewrite to optimize `dot(kron(a, b), c)`

Open
#1,043 0 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

graph rewriting linalg performance
Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

Description

This can be re-written according to the following relationship:

$$
(A \otimes B) C = \text{vec}(B X A^T)
$$

Where $\otimes$ is the kronecker product, and the $\text{vec}$ operation ravels a matrix in column-major order. $X$ is a matrix formed by reshaping $C$ (in column-major order) to conform with $B$. This avoids working with the large kronecker product matrix, and instead gets the result in terms of the much smaller components. Code example:

n = 100
a, b = np.random.normal(size=(2, n, n))
c = np.random.normal(size=(n ** 2, ))

def kronAB_C_clever(a, b, c):
    return (b @ c.reshape((n, n)).T @ a.T).T.ravel()

def direct(a, b, c):
    K = np.kron(a, b)
    x2 = K @ c

%timeit kronAB_C_clever(a, b, c)
73.5 μs ± 3.61 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

%timeit direct(a, b, c)
245 ms ± 72.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

This trick is already used in PyMC here, but only in a limited context. PyMC applies this identity to solve_triangular as well, but it can (and should) also be applied to other types of solve.

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

Start by locating PyTensor's rewrite rules for dot and kron, then compare them with the limited solve_triangular use described in the issue and the linked PyMC math.py implementation. Validate the Kronecker identity and determine which solve operations should use it; done means the applicable rewrites avoid materializing the large Kronecker product and preserve the expected result.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.