Einsums are the regexes of tensor programming. Should be avoided at all costs IMO. Ideally we should be able to write native loops that get auto-vectorized into einsums for which there is a CUDA/PTX emitting factory. But for some reason neither PyTorch nor JAX/TF took this route and now we are here.
Some of the einsum expressions I have seen for grouped multi headed/query attention is mind-boggling and they get shipped to prod.
Some of the einsum expressions I have seen for grouped multi headed/query attention is mind-boggling and they get shipped to prod.