search

theislab / cpa

The Compositional Perturbation Autoencoder (CPA) is a deep generative framework to learn effects of perturbations at the single-cell level. CPA performs OOD predictions of unseen combinations of drugs, learns interpretable embeddings, estimates dose-response curves, and provides uncertainty estimates.
BSD 3-Clause "New" or "Revised" License
76 stars 17 forks source link

fix jax deps #43

Closed zgr2788 closed 3 months ago

zgr2788 commented 4 months ago

Issue

Encountered by @erialc-cal

Importing cpa fails after following the installation tutorial due to the call to jax.random.KeyArray in ... which was already depracated in 0.4.16 and recently (Feb. 6th) removed in 0.4.24, see here for more details.

>>> import cpa [rank: 0] Global seed set to 0 ...[trace]... AttributeError: module 'jax.random' has no attribute 'KeyArray'

How to Reproduce

  1. conda create -n cpa python=3.9
  2. conda activate cpa
  3. pip install git+https://github.com/theislab/cpa
  4. python
  5. >>> import cpa

Fix/Changes

Added jax version upper bound in pyproject.toml, also added same for jaxlib to avoid consistency errors for jax.

ArianAmani commented 3 months ago

Thank you for noticing the issue, I am merging your PR to the main branch.