AttributeError: module 'jax.tree_util' has no attribute 'register_pytree_with_keys_class' about whisper-jax

[talipturkmen]

you can upgrade jax in colab using
!pip install --upgrade jax jaxlib
it works for me after upgrade

[silvacarl2]

LOL:

RuntimeError:
As of JAX 0.4.0, JAX only supports TPU VMs, not the older Colab TPUs.

We recommend trying Kaggle Notebooks
(https://www.kaggle.com/code, click on "New Notebook" near the top) which offer
TPU VMs. You have to create an account, log in, and verify your account to get
accelerator support.
Once you do that, there's a new "TPU 1VM v3-8" accelerator option. This gives
you a TPU notebook environment similar to Colab, but using the newer TPU VM
architecture. This should be a less buggy, more performant, and overall better
experience than the older TPU node architecture.

It is also possible to use Colab together with a self-hosted Jupyter kernel
running on a Cloud TPU VM. See
https://research.google.com/colaboratory/local-runtimes.html
for details.

[silvacarl2]

yes this is going to be rough:

Cloud TPU v3: $8.00 hour
Cloud TPU v4: $12.88 per hour

[silvacarl2]

by the way, it works perfect on Kaggle!

[silvacarl2]

this is also depressing:

Creating TPU node "whisper-tpu-test-server" failed. Error: Request failed with unknown error

[silvacarl2]

get this same error on TPU V2.XX

AttributeError: module 'jax.tree_util' has no attribute 'register_pytree_with_keys_class'

Except for Kallg,e all TPU regions are maxed out for TPU V3.xx

[sanchit-gandhi]

Hey @silvacarl2 - it's super sad that Google Colab is no longer really support its TPU architectures. I also initially tried writing the notebook in Colab, but got the same recommendation as you and so switched to Kaggle notebooks.

There is indeed a bit of a queue to get a TPU VM v3-8 with Kaggle, but once you do it's worth it (20 hours free per week is very generous).

[sanchit-gandhi]

Also preemptible TPUs will be cheaper if you're happy to accept the risk that you might lose a TPU in return for a lower cost

[silvacarl2]

yes, this works great on Kaggle.

we wanted to test it out and benchmark it but we could not eyt get it to run on anything except Kaggle.

[silvacarl2]

no worries. this is awesome work because it FLIES on Kaggle.