einops
torch>=1.6
triton
