torch
einops
matplotlib
jaxtyping
