JAX connection with Google AI Platform Notebook

[humtumiit]

I have created a AI Platform Notebook hosted on VM in google cloud and a TPU node. I would like to use JAX as backend which should use TPU hardware to perform complex operations but unfortunately I am not able to do this.

When I see the list of TPU by running the command in notebook,
!gcloud compute tpus list --zone=us-central1-c

I get the following
NAME ZONE ACCELERATOR_TYPE NETWORK RANGE STATUS
gptj6btpu us-central1-c v2-8 default 10.22.214.0/29 READY
.0

which shows TPU is up and running and accessible from the notebook,

I installed and imported JAX and want to see the number of devlies using

print(jax.devices())

But I get only one CPU instance and message that No GPU/TPU has been detected and falling it down to CPU.

Can anyone suggest, what I am doing wrong and how to detect TPU from the AI Platform hosted notebook?