Jax dependency error, solution other than downgrade? about scvi-tools HOT 6 CLOSED

scvi-tools                1.1.2                    pypi_0    pypi

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[1], line 1
----> 1 from flax.training import train_state

File ~/miniconda3/envs/test_env/lib/python3.10/site-packages/flax/training/train_state.py:17
      1 # Copyright 2024 The Flax Authors.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     12 # See the License for the specific language governing permissions and
     13 # limitations under the License.
     15 from typing import Any, Callable
---> 17 import optax
     19 from flax import core, struct
     20 from flax.linen.fp8_ops import OVERWRITE_WITH_GRADIENT

File ~/miniconda3/envs/test_env/lib/python3.10/site-packages/optax/__init__.py:17
      1 # Copyright 2019 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 """Optax: composable gradient processing and optimization, in JAX."""
---> 17 from optax import contrib
     18 from optax import losses
     19 from optax import monte_carlo

File ~/miniconda3/envs/test_env/lib/python3.10/site-packages/optax/contrib/__init__.py:17
      1 # Copyright 2021 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 """Contributed optimizers in Optax."""
---> 17 from optax.contrib.cocob import cocob
     18 from optax.contrib.cocob import COCOBState
     19 from optax.contrib.complex_valued import split_real_and_imaginary

File ~/miniconda3/envs/test_env/lib/python3.10/site-packages/optax/contrib/cocob.py:25
     23 import jax.numpy as jnp
     24 import jax.tree_util as jtu
---> 25 from optax._src import base
     28 class COCOBState(NamedTuple):
     29   """State for COntinuous COin Betting."""

File ~/miniconda3/envs/test_env/lib/python3.10/site-packages/optax/_src/base.py:19
     15 """Base interfaces and datatypes."""
     17 from typing import Any, Callable, NamedTuple, Optional, Protocol, runtime_checkable, Sequence, Union
---> 19 import chex
     20 import jax
     21 import jax.numpy as jnp

File ~/miniconda3/envs/test_env/lib/python3.10/site-packages/chex/__init__.py:17
      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 ~/miniconda3/envs/test_env/lib/python3.10/site-packages/chex/_src/asserts.py:26
     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 ~/miniconda3/envs/test_env/lib/python3.10/site-packages/chex/_src/asserts_internal.py:34
     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 ~/miniconda3/envs/test_env/lib/python3.10/site-packages/chex/_src/pytypes.py:54
     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 ~/miniconda3/envs/test_env/lib/python3.10/site-packages/jax/_src/deprecations.py:54, in deprecation_getattr.<locals>.getattr(name)
     52   warnings.warn(message, DeprecationWarning, stacklevel=2)
     53   return fn
---> 54 raise AttributeError(f"module {module!r} has no attribute {name!r}")

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

flax                      0.8.2 
jax                       0.4.26
jaxlib                    0.4.26
chex                      0.1.7
optax                     0.2.1

scvi-tools 1.1.2 requires flax, which is not installed.
Successfully installed jax-0.4.26 jaxlib-0.4.26+cuda12.cudnn89 nvidia-cuda-nvcc-cu12-12.4.131