what to do:
- check if the Newton-Schultz and the KFAC are available in python (even in triton or whatever)
- start with one kernel for each big part, then fusion
- KFAC approximation: KFAC inverse * gradient. curvlinops is one option, backpack is another, test it.
- sinkhorn approx.: POT python package
- N-S iteration: lowest hanging fruit kábé
- multi-scale loss function