Butterfly matrix multiplication in PyTorch - View it on GitHub