numpy torch torchvision transformers scipy ml_collections datasets matplotlib scikit-learn einops