torch
scikit-learn
einops
