numpy>=1.20
jax
jaxlib
typing_extensions
