typing_extensions>=4.5
array_api_compat
array_api_extra
numpy

[jax]
jax

[torch]
torch
