Set up JAX ERROR

[gitbim]

The TPU Version wont start.

ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
tensorflow 2.12.0 requires numpy<1.24,>=1.22, but you have numpy 1.25.1 which is incompatible.
Successfully installed jax-0.4.13 jaxlib-0.4.13
WARNING: The following packages were previously imported in this runtime:
  [jax,jaxlib]
You must restart the runtime in order to use newly installed versions.

I can set up JAX with "jax==0.3.25 jaxlib==0.3.25", but then interface and Walk! wont work. Could you have a look at your TPU version with you basic settings?

[aduchon]

This order worked for me in a colab. You basically have to turn time back to November, 2022, before the jax0.4 versions. I did not even have to restart like it asks.

!pip install numpy==1.23.5
!pip install chex==0.1.6
!pip install flax==0.6.0
!pip install orbax-checkpoint==0.1.1
!pip install diffusers==0.9.0 transformers==4.24.0 ftfy==6.1.1
!pip install jax==0.3.25 jaxlib==0.3.25
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu('tpu_driver_20221011')
jax.devices()

[aduchon]

Note that the jax/tpu version can not deal with negative prompts. I ended up going back to the regular gpu version, which also seemed faster.

[nateraw]

Thank you @aduchon for the snippet! I'll look into this and update the colab accordingly.