Which jax version to use?

[resurgo97]

I installed jax with the below command,

pip install --upgrade "jax[cpu]" install -q neural_tangents==0.2.2 livelossplot

but I come across the error jax.api does not exist while importing neural tangent package and searched to find that I need to upgrade neural tangent package, or downgrade jax.

Both worked, but in either case I run into another kind of error while executing the other codes.

I also tried running the codes on Colab, by clicking on the badge in the notebook file, but also had the same error: jax.api does not exist.

I find it hard to find any proper combination of those two packages; neural tangent and jax.

Could anyone help me please?