Running Colab or building manually - both broken?

[vican9000]

I've tried running the provided Google Colab notebook, without any luck (Colab and TPU versioning issues).
After trying to build manually, I keep getting dependency problems with t5x, flax and jax (tried many combinations).
Is this completely broken or there might be a way to run it somehow? Thanks!

[hcps-pressleaf]

Still seems broken out of the box. jax imports aren't compatible with the old version of t5x, but updating it messes with imports that happen within the model. Help!