annotation notebook does not run within container

The recommended interactive container built nicely on ubuntu 22.04 but attempt to run first chunk of annotation.ipynb led to

2024-02-13 14:15:57 - INFO - generated new fontManager

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[1], line 8
      6 import pandas as pd
      7 import scanpy as sc
----> 8 import scvi
      9 import numpy as np
     11 scvi.settings.seed = 0

File [/usr/local/lib/python3.11/site-packages/scvi/__init__.py:11](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/scvi/__init__.py#line=10)
      8 from ._settings import settings
     10 # this import needs to come after prior imports to prevent circular import
---> 11 from . import autotune, data, model, external, utils, criticism
     13 from importlib.metadata import version
     15 package_name = "scvi-tools"

File [/usr/local/lib/python3.11/site-packages/scvi/autotune/__init__.py:1](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/scvi/autotune/__init__.py#line=0)
----> 1 from ._manager import TuneAnalysis, TunerManager
      2 from ._tuner import ModelTuner
      3 from ._types import Tunable, TunableMixin

File [/usr/local/lib/python3.11/site-packages/scvi/autotune/_manager.py:11](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/scvi/autotune/_manager.py#line=10)
      9 import lightning.pytorch as pl
     10 import rich
---> 11 from chex import dataclass
     13 try:
     14     import ray

File [/usr/local/lib/python3.11/site-packages/chex/__init__.py:17](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/chex/__init__.py#line=16)
      1 # Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     13 # limitations under the License.
     14 # ==============================================================================
     15 """Chex: Testing made fun, in JAX!"""
---> 17 from chex._src.asserts import assert_axis_dimension
     18 from chex._src.asserts import assert_axis_dimension_comparator
     19 from chex._src.asserts import assert_axis_dimension_gt

File [/usr/local/lib/python3.11/site-packages/chex/_src/asserts.py:26](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/chex/_src/asserts.py#line=25)
     23 import unittest
     24 from unittest import mock
---> 26 from chex._src import asserts_internal as _ai
     27 from chex._src import pytypes
     28 import jax

File [/usr/local/lib/python3.11/site-packages/chex/_src/asserts_internal.py:34](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/chex/_src/asserts_internal.py#line=33)
     31 from typing import Any, Sequence, Union, Callable, List, Optional, Set, Tuple, Type
     33 from absl import logging
---> 34 from chex._src import pytypes
     35 import jax
     36 from jax.experimental import checkify

File [/usr/local/lib/python3.11/site-packages/chex/_src/pytypes.py:54](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/chex/_src/pytypes.py#line=53)
     52 Numeric = Union[Array, Scalar]
     53 Shape = jax.core.Shape
---> 54 PRNGKey = jax.random.KeyArray
     55 PyTreeDef = jax.tree_util.PyTreeDef
     56 Device = jax.Device

File [/usr/local/lib/python3.11/site-packages/jax/_src/deprecations.py:53](http://127.0.0.1:8888/usr/local/lib/python3.11/site-packages/jax/_src/deprecations.py#line=52), in deprecation_getattr.<locals>.getattr(name)
     51   warnings.warn(message, DeprecationWarning, stacklevel=2)
     52   return fn
---> 53 raise AttributeError(f"module {module!r} has no attribute {name!r}")

AttributeError: module 'jax.random' has no attribute 'KeyArray'

Start with annotation.ipynb and reproduce the first chunk in the recommended container, focusing on the scvi import and the reported JAX/chex compatibility error. Inspect the installed dependency versions and verify that the notebook's initial imports run successfully without the AttributeError.