numpy
torch
einops
ray[train]
