chex
einops
numpy
jax
haiku