OWL-ViT inference playround: ModuleNotFoundError: No module named 'jax.experimental.gda_serialization'

I tried to run the sample notebook (https://colab.research.google.com/github/google-research/scenic/blob/main/scenic/projects/owl_vit/notebooks/OWL_ViT_inference_playground.ipynb), and encountered the following error while running the second cell:

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-2-732db45ee9f0>](https://localhost:8080/#) in <cell line: 10>()
      8 import numpy as np
      9 from scenic.projects.owl_vit import configs
---> 10 from scenic.projects.owl_vit import models
     11 
     12 from scenic.projects.owl_vit.notebooks import inference

5 frames
[/usr/local/lib/python3.10/dist-packages/orbax/checkpoint/type_handlers.py](https://localhost:8080/#) in <module>
     22 from etils import epath
     23 import jax
---> 24 from jax.experimental.gda_serialization import serialization
     25 from jax.experimental.gda_serialization.serialization import get_tensorstore_spec
     26 import jax.numpy as jnp

ModuleNotFoundError: No module named 'jax.experimental.gda_serialization'

---------------------------------------------------------------------------
NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.
---------------------------------------------------------------------------

I think this is caused by version incompatibility of jax. Could you tell me the appropriate versions of dependencies?

Thank you in advance.

P.S. I found the error above when I used TPU as run time. After changing it to GPU (T4), I encountered another error in the second cell:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-2-732db45ee9f0>](https://localhost:8080/#) in <cell line: 10>()
      8 import numpy as np
      9 from scenic.projects.owl_vit import configs
---> 10 from scenic.projects.owl_vit import models
     11 
     12 from scenic.projects.owl_vit.notebooks import inference

14 frames
[/usr/local/lib/python3.10/dist-packages/numpy/testing/_private/utils.py](https://localhost:8080/#) in <module>
     55 IS_PYSTON = hasattr(sys, "pyston_version_info")
     56 HAS_REFCOUNT = getattr(sys, 'getrefcount', None) is not None and not IS_PYSTON
---> 57 HAS_LAPACK64 = numpy.linalg._umath_linalg._ilp64
     58 
     59 _OLD_PROMOTION = lambda: np._get_promotion_state() == 'legacy'

AttributeError: module 'numpy.linalg._umath_linalg' has no attribute '_ilp64'

google-research / scenic

OWL-ViT inference playround: ModuleNotFoundError: No module named 'jax.experimental.gda_serialization' #957