Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
jessegrabowski opened this issue Oct 20, 2024 · 0 comments · May be fixed by #1090
Open

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

jessegrabowski opened this issue Oct 20, 2024 · 0 comments · May be fixed by #1090

Comments

@jessegrabowski
Copy link
Member

jessegrabowski commented Oct 20, 2024

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.

@tanish1729 tanish1729 linked a pull request Nov 14, 2024 that will close this issue
10 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants