Skip to content

Instantly share code, notes, and snippets.

View yuxuan-z19's full-sized avatar
💤
lying flat

yuxuan-z yuxuan-z19

💤
lying flat
  • Beijing, China
View GitHub Profile
@norabelrose
norabelrose / kronecker_decompose.py
Last active May 1, 2025 08:46
Fast, optimal Kronecker decomposition
from einops import rearrange
from torch import Tensor
import torch
def kronecker_decompose(
A: Tensor, m: int, n: int, *, k: int = 1, niter: int = 10
) -> tuple[Tensor, Tensor]:
"""Frobenius-optimal decomposition of `A` into a sum of `k` Kronecker products.