pymc-devs / pymc-devs/pytensor
Add rewrite to optimize `dot(kron(a, b), c)`
Nobody has claimed this yet.
- 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
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
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