flax==0.1.0
numpy==1.18.4
tensorflow-cpu==2.2.0
tensorflow-datasets==3.1.0
tensorflow-probability==0.9.0
