Installation jax related packages errs due to deprecations

[Hrovatin]

I tried installing a fresh scvi env via conda, as per install guide, but got bellow wee on scvi-tools import

>>> from scvi.model import SCVI

/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/flax/struct.py:132: FutureWarning: jax.tree_util.register_keypaths is deprecated, and will be removed in a future release. Please use `register_pytree_with_keys()` instead.
  jax.tree_util.register_keypaths(data_clz, keypaths)
/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/flax/struct.py:132: FutureWarning: jax.tree_util.register_keypaths is deprecated, and will be removed in a future release. Please use `register_pytree_with_keys()` instead.
  jax.tree_util.register_keypaths(data_clz, keypaths)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/__init__.py", line 11, in <module>
    from . import data, model, external, utils
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/data/__init__.py", line 25, in <module>
    from ._preprocessing import (
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/data/_preprocessing.py", line 12, in <module>
    from scvi.model._utils import parse_device_args
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/model/__init__.py", line 2, in <module>
    from ._amortizedlda import AmortizedLDA
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/model/_amortizedlda.py", line 15, in <module>
    from scvi.module import AmortizedLDAPyroModule
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/module/__init__.py", line 1, in <module>
    from ._amortizedlda import AmortizedLDAPyroModule
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/module/_amortizedlda.py", line 15, in <module>
    from scvi.module.base import PyroBaseModuleClass, auto_move_data
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/module/base/__init__.py", line 1, in <module>
    from ._base_module import (
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/scvi/module/base/_base_module.py", line 14, in <module>
    from flax.training import train_state
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/flax/training/train_state.py", line 19, in <module>
    import optax
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/optax/__init__.py", line 18, in <module>
    from optax import losses
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/optax/losses/__init__.py", line 17, in <module>
    from optax.losses._classification import convex_kl_divergence
  File "/Users/karin.hrovatin/miniconda3/envs/scvi-test/lib/python3.9/site-packages/optax/losses/_classification.py", line 64, in <module>
    chex.warn_deprecated_function,
AttributeError: module 'chex' has no attribute 'warn_deprecated_function'

Fixed by (at least SCVI import):

• pip install jaxlib==0.4.13
• pip install jax==0.4.13
• pip install numpyro==0.12.1
• pip install chex==0.1.82
• pip install optax==0.2.0

The downgrades worked as there were deprecations not accounted for in scvi-tools

[canergen]

Hi. Chex is no requirement anymore (is my understanding). Can you print the version of scVI-tools? You want to update to Python>3.9 to install the most recent release.

[Hrovatin]

The installation guide has python 3.9 fixed: https://docs.scvi-tools.org/en/1.0.4/installation.html#conda After creating the env I followed the tutorial to install from conda-forge

scvi-tools 1.1.6.post2 pyhd8ed1ab_1 conda-forge

[canergen]

You should look at the latest installation guide (set to 3.12): https://docs.scvi-tools.org/en/latest/installation.html. I looked again at your message and the issue is within optax. It's not on our end to fix it but it's raised in optax https://github.com/google-deepmind/optax/issues/995 and it is their thing to fix it (apparently requirements are ignored by conda).