Issues with Requirements

[Yanay1]

Hello,

What are the exact versions of jax and jaxlib that are required? I have tried with a variety of jax and jaxlib versions including those that are not available from pypi and get a variety of different import errors including:

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

or

jaxlib.xla_client' has no attribute 'ComputationBuilder

or

AttributeError: module 'jax.tree_util' has no attribute 'register_pytree_with_keys_class'

This happens either when using the requirements.txt or the poetry installation.

[subercui]

Hi, we encourage the poetry installation for now. We are trying to push out a working pypi release in the coming days, which we expect to be an easier installation way.

In the meantime, we don't explicitly require jaxlib. So I figure this issue is probably related to a specific dependency. Can you please upload the full error messages, so I can take a further look at it?

[Yanay1]

I am also having issues with the flash attention installation- I'll wait for the pip release to try again. Thanks!

[subercui]

Sure, and feel free to reopen if you still encounter issues later.

[ApolloLSY]

I encountered the same issue about the jax version.

Though poetry is kindly ready now, the 0.3.10 version set in poetry currently cannot be downloaded. Poetry reports an error saying there is no 0.3.10 version of the jaxlib package.

I found out this is because some of the very old versions have been deleted officially and cannot be downloaded from pip. However, I seek to install it like this: pip install --upgrade "jax[cuda12_pip]==0.3.13" "jaxlib==0.3.10" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html.

But after installation, when running it, version 0.3.10 has an error 'jaxlib.mlir._mlir_libs._mlir.ir' has no attribute '_BaseContext').

I found a report on this issue on the jax GitHub page (jax-ml/jax#12526), where developers mentioned that an newer version is needed.

So I try not to stick with the recommended versions 0.3.10 and just try the latest 0.4.3, but 0.4.3 had problems tooModuleNotFoundError: No module named 'jax.extend'. I then tried 0.3.24, and it still had problems module 'jax' has no attribute 'Device'. I gave up.

But then I found jax is a dependency package for scvi, and found

You can just deleting all the usages of scvi in fine-tuning

So, to anyone reading this, if you encounter the same problems while training your own data, try this.

[LIHAO-0517]

placeholder