Can confirm. I get the exact same error

will take a look thx for reporting. 😄

Can confirm its not running as it was previously. I got it "running" by restarting runtime after all the installs, but it was very very slow. There was a new release of jax recently (Feb 16th), so I assume that's causing it. Maybe need to either:

• downgrade jax or
• perhaps just need to change the tpu_driver_20221011 value we used in the notebook to a newer one.

I have no clue. @charlielito you mind having a look at this?

I got this to work last night by doing the following:

• Just using jax that comes with colab (version 0.3.25)
• changing setup to not specify build number, since they update it with the release jax.tools.colab_tpu.setup_tpu()

Still would prefer to wait for @charlielito to weigh in before I make any changes to colab

Thanks this seems to work indeed!

I got this to work last night by doing the following:

• Just using jax that comes with colab (version 0.3.25)
• changing setup to not specify build number, since they update it with the release jax.tools.colab_tpu.setup_tpu()

Still would prefer to wait for @charlielito to weigh in before I make any changes to colab

I was on vacation :p I'll take a look this week

All good @charlielito I hope you enjoyed your time off ❤️ . No rush here - let me know if I should merge my hotfix or if you think there's a better way whenever you get the chance.

For me it doesn't work.
I get:
RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.4.

For me it doesn't work. I get: RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.4.

Yeah, the new release of jax/jaxlib/flax is not compatible with the TPU API we are using. I'll work on the versions needed to run it again