Skip to content

Kronecker linalg rewrites#2179

Open
jessegrabowski wants to merge 5 commits into
pymc-devs:mainfrom
jessegrabowski:kron-solve
Open

Kronecker linalg rewrites#2179
jessegrabowski wants to merge 5 commits into
pymc-devs:mainfrom
jessegrabowski:kron-solve

Conversation

@jessegrabowski
Copy link
Copy Markdown
Member

Adds several rewrites related to kronecker products that avoid materializing the full big matrices:

  • solve_of_kron (solvers.py) — push solve through a Kronecker product:

    solve(kron(A, B), b)     ->  vec(solve(B, solve(A, unvec(b)).T).T)
    solve(kron(A, B).T, b)   ->  same, with A->A.T, B->B.T   (transpose peel)
    
  • dot_of_kron (products.py) — push matmul through a Kronecker product (dual identity):

    kron(A, B) @ X     ->  reshape(A @ reshape(X, (m,p,k)) @ B.T, (m*p, k))
    kron(A, B).T @ X   ->  same, with A->A.T, B->B.T
    
  • solve_of_kron_plus_diag_noise (solvers.py):

    solve(kron(K1, K2) + sigma**2*I, b)
        ->  Q @ ((Q.T @ b) / d)
            where  Q = kron(Q1, Q2),  d = kron_diag(d1, d2) + sigma**2
    
  • det_of_kron_plus_diag_noise (summary.py):

    det(kron(K1, K2) + sigma**2*I)
        ->  prod(kron_diag(d1, d2) + sigma**2)
    

The prupose is to eliminate completely the KroneckerNormal distribution in PyMC, users will just get the best possible logp if they put a kronecker covariance matrix into a normal. solve and dot of kron are generally useful, but the ones with diagonal noise are pretty specialized, so I had them track KroneckerProduct itself, rather than Solve or Dot.

@jessegrabowski jessegrabowski added enhancement New feature or request graph rewriting linalg Linear algebra labels May 28, 2026
@jessegrabowski jessegrabowski changed the title Kronecker solve rewrites Kronecker linalg rewrites May 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request graph rewriting linalg Linear algebra

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant