jax==0.2.21
numpy==1.19.5
scipy==1.7.1
tensorflow_datasets==4.4.0
