Error running OWL-ViT v2 example colab

[samuela]

I just tried running this example colab, but am seeing an error:

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-2-6cf721b232cb>](https://sal4s0c6ilt-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20230926-060131_RC00_568496853#) in <cell line: 9>()
      7 from matplotlib import pyplot as plt
      8 import numpy as np
----> 9 from scenic.projects.owl_vit import models
     10 from scenic.projects.owl_vit.configs import clip_b16 as config_module
     11 from scenic.projects.owl_vit.notebooks import inference

5 frames
[/usr/local/lib/python3.10/dist-packages/orbax/checkpoint/type_handlers.py](https://sal4s0c6ilt-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20230926-060131_RC00_568496853#) 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.
---------------------------------------------------------------------------

[mjlm]

This may be due to a jax version issue. Could you try again in a freshly installed environment and let me know how it goes?

When I test it in a fresh install, the models import works. I then hit #859, which we'll fix soon (see that issue for a workaround).