numpy
torch
functorch
flax
jax
