pymc-devs / pymc-devs/pytensor
Rewrite `Solve` involving `Kron`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
Define $\text{vec}(X)$ as the ravel operator that works column-wise (that's the math convention, so in code it's x.reshape(-1, order="F").
Given $(A \otimes B)x = y$. If $y = \text{vec}(Y)$, then $x = \text{vec}(B^{-1}YA^{-T})$
This identity suggests a rewrite for solve Ops involving Kron:
import pytensor.tensor as pt
import pytensor
import numpy as np
A, B = pt.dmatrices('A', 'B')
y = pt.dvector('y')
n, m = A.shape[0], B.shape[0]
x1 = pt.linalg.solve(pt.linalg.kron(A, B), y)
x2 = pt.linalg.solve(A,
pt.linalg.solve(B, y.reshape((n, m)).T).T).ravel()
fn1 = pytensor.function([A, B, y], x1)
fn2 = pytensor.function([A, B, y], x2)
# Show equivalence
rng = np.random.default_rng()
n = 50
a_val, b_val = rng.normal(size=(2, n, n))
y_val = rng.normal(size=(n * n))
np.allclose(fn1(a_val, b_val, y_val), fn2(a_val, b_val, y_val)) # True
Timings:
%timeit fn1(a_val, b_val, y_val) # 83.9 ms ± 626 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit fn2(a_val, b_val, y_val) # 178 μs ± 8.41 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
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 tracing the Solve and Kron Ops and the rewrite machinery that handles them; the issue names no files or tests. Use the provided fn1 and fn2 expressions to verify equivalent results, then confirm that the rewrite avoids explicitly solving the Kronecker-product system and preserves the expected behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- backend-api-design, performance
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100