Add cuda accelerated jax for scvi-tools (#1028)

* Add cuda accelerated jax for scvi-tools

* Revert removal of comment

* Update CHANGELOG [ci skip]

Author: DriesSchaumont

CHANGELOG.md

@@ -20,6 +20,8 @@
 
 * Fix running `neigbors_leiden_umap` workflow with `-stub` enabled (PR #1026).
 
+* Add missing CUDA enabled `jaxlib` to components that use `scvi-tools` (`scanvi`, `scarches`, `scvi` and `totalvi`) (PR #1028)
+
 # openpipelines 2.1.0
 
 ## BREAKING CHANGES

src/annotate/scanvi/config.vsh.yaml

@@ -146,6 +146,7 @@ engines:
       __merge__: [/src/base/requirements/anndata_mudata.yaml, /src/base/requirements/scanpy.yaml, .]
     - type: python
       packages:
+        - jax[cuda]
         - scvi-tools~=1.1.5
   test_setup:
     - type: python

src/integrate/scarches/config.vsh.yaml

@@ -147,6 +147,7 @@ engines:
     - type: python
       __merge__: [/src/base/requirements/anndata_mudata.yaml, .]
       packages:
+        - jax[cuda]
         - scvi-tools~=1.1.5
   test_setup:
     - type: python

src/integrate/scvi/config.vsh.yaml

@@ -233,6 +233,7 @@ engines:
       __merge__: [/src/base/requirements/anndata_mudata.yaml, /src/base/requirements/scanpy.yaml, .]
     - type: python
       packages:
+        - jax[cuda]
         - scvi-tools~=1.1.5
   test_setup:
     - type: python

src/integrate/totalvi/config.vsh.yaml

@@ -120,6 +120,7 @@ engines:
     - type: python
       __merge__: [/src/base/requirements/anndata_mudata.yaml, .]
       packages:
+        - jax[cuda]
         - scvi-tools~=1.1.5
   test_setup:
     - type: python