torch
einops
datasets
dataclasses-json
jaxtyping
