AttributeError: 'jaxlib.tpu_client_extension.PyTpuBuffer' object has no attribute 'dtype'

[haroonmoh]

/usr/local/lib/python3.7/dist-packages/jax/interpreters/pxla.py in make_sharded_device_array(aval, sharding_spec, device_buffers, indices)
    480 
    481   for buf in device_buffers:
--> 482     buf.aval = ShapedArray(buf.shape, buf.dtype)
    483 
    484   if (_USE_CPP_SDA and

I keep getting this error. I am using Trax 1.3.9, Jax 0.2.20, And Jaxlib 0.1.71.
I am using a TPU on Google Colab. My code for using the TPU is below. I'm not sure if I am doing something wrong.

if 'TPU_DRIVER_MODE' not in globals():
  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver_nightly'
  resp = requests.post(url)
  TPU_DRIVER_MODE = 1

from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
print(config.FLAGS.jax_backend_target)

[syzymon]

Also reproduced while using Trax 1.3.9 on a colab TPU. As a temporary workaround, downgrading jax/jaxlib to the following helps:

!pip uninstall -yq jaxlib
!pip install jaxlib==0.1.64+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
!pip uninstall -yq jax
!pip install jax==0.2.12

[hawkinsp]

I'm pretty sure this problem was fixed by jax 0.2.21, which was released yesterday. Please try it out?

[syzymon]

Updating to 0.2.21 seems to resolve the issue indeed, thanks! Also attaching jaxlib version for completeness:

jax                           0.2.21
jaxlib                        0.1.71+cuda111