JAX issue

[anhdnguy]

I was able to install scvi-tools via conda just fine, but when I import the module, it returns this error

cannot import name 'ShapedArray' from 'jax'

Versions:

scvi-tools = 1.0.2
jax = 0.4.14

[martinkim0]

Hi, this issue is originating from flax since it looks like the package is not being updated correctly on conda (latest on conda is 0.6.1 but latest on pypi is 0.7.0). For now, this can be fixed by installing via pip:

pip install -U scvi-tools

and we'll find a fix for the conda installation in our next patch (1.0.3). Hope this helps!

[anhdnguy]

@martinkim0 Thanks. I installed the old version and it worked. I will look out for the new release

[maarten-devries]

I also needed pip install -U flax