jax>=0.4.4 dependency

[gokceneraslan]

Comment:

Any particular reason of why this package depends on jax>=0.4.4 while the original scvi-tools doesn't have a specific version dependency?

Also, although I have jax 0.4.28 installed, I'm still getting an error when I try to install latest scvi, any ideas?

(scvi-env) :/data$ mamba install scvi-tools==1.1.5

                  __    __    __    __
                 /  \  /  \  /  \  /  \
                /    \/    \/    \/    \
███████████████/  /██/  /██/  /██/  /████████████████████████
              /  / \   / \   / \   / \  \____
             /  /   \_/   \_/   \_/   \    o \__,
            / _/                       \_____/  `
            |/
        ███╗   ███╗ █████╗ ███╗   ███╗██████╗  █████╗
        ████╗ ████║██╔══██╗████╗ ████║██╔══██╗██╔══██╗
        ██╔████╔██║███████║██╔████╔██║██████╔╝███████║
        ██║╚██╔╝██║██╔══██║██║╚██╔╝██║██╔══██╗██╔══██║
        ██║ ╚═╝ ██║██║  ██║██║ ╚═╝ ██║██████╔╝██║  ██║
        ╚═╝     ╚═╝╚═╝  ╚═╝╚═╝     ╚═╝╚═════╝ ╚═╝  ╚═╝

        mamba (1.2.0) supported by @QuantStack

        GitHub:  https://github.com/mamba-org/mamba
        Twitter: https://twitter.com/QuantStack

█████████████████████████████████████████████████████████████

Looking for: ['scvi-tools==1.1.5']

conda-forge/linux-64                                        Using cache
conda-forge/noarch                                          Using cache
bioconda/linux-64                                           Using cache
bioconda/noarch                                             Using cache
pkgs/r/noarch                                                 No change
pkgs/r/linux-64                                               No change
pkgs/main/linux-64                                            No change
pkgs/main/noarch                                              No change

Pinned packages:
  - python 3.9.*

Could not solve for environment specs
Encountered problems while solving:
  - package scvi-tools-1.1.5-pyhd8ed1ab_0 requires jax >=0.4.4, but none of the providers can be installed

The environment can't be solved, aborting the operation

(scvi-env) :/data$ mamba list jax
# packages in environment at /home/user/.miniconda3/envs/scvi-env:
#
# Name                    Version                   Build  Channel
jax                       0.4.28             pyhd8ed1ab_0    conda-forge
jaxlib                    0.4.28          cpu_py39h32154bb_0    conda-forge

[canergen]

Hi, the error looks like a weird conda environment solving issue. I assume 0.4.4-0.4.9 would work? Unfortunately, JAX has a high tendency of creating patch fixes that are not backward compatible and sometimes even doesn't handle those with deprecation warnings. I assume the drop of jax.experimental in 0.4.4 created the lower version requirement.