From 389cb708c617e73b874fc05a8d03a1a6d9ebc658 Mon Sep 17 00:00:00 2001
From: Dries Schaumont <5946712+DriesSchaumont@users.noreply.github.com>
Date: Wed, 14 May 2025 18:55:58 +0000
Subject: [PATCH 1/3] Add cuda accelerated jax for scvi-tools

---
 src/annotate/scanvi/config.vsh.yaml    | 1 +
 src/integrate/scarches/config.vsh.yaml | 1 +
 src/integrate/scvi/config.vsh.yaml     | 3 ++-
 src/integrate/totalvi/config.vsh.yaml  | 1 +
 4 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/annotate/scanvi/config.vsh.yaml b/src/annotate/scanvi/config.vsh.yaml
index 27bc8107464..96b22e15779 100644
--- a/src/annotate/scanvi/config.vsh.yaml
+++ b/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
diff --git a/src/integrate/scarches/config.vsh.yaml b/src/integrate/scarches/config.vsh.yaml
index 1460ef2847b..d5c6d2953d1 100644
--- a/src/integrate/scarches/config.vsh.yaml
+++ b/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
diff --git a/src/integrate/scvi/config.vsh.yaml b/src/integrate/scvi/config.vsh.yaml
index 7bd5ffa0ec5..b0adcc27443 100644
--- a/src/integrate/scvi/config.vsh.yaml
+++ b/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
@@ -240,7 +241,7 @@ engines:
 
 runners:
 - type: executable
-  # docker_run_args: ["--gpus all"]
+  docker_run_args: ["--gpus all"]
 - type: nextflow
   directives:
     label: [midcpu, midmem, gpu, highdisk]
diff --git a/src/integrate/totalvi/config.vsh.yaml b/src/integrate/totalvi/config.vsh.yaml
index c733ab682ba..9a150326ef0 100644
--- a/src/integrate/totalvi/config.vsh.yaml
+++ b/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

From 3412b8d1d56a6ffbb6d76f99f38c14e71ff6ccd1 Mon Sep 17 00:00:00 2001
From: Dries Schaumont <5946712+DriesSchaumont@users.noreply.github.com>
Date: Wed, 14 May 2025 18:57:34 +0000
Subject: [PATCH 2/3] Revert removal of comment

---
 src/integrate/scvi/config.vsh.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/integrate/scvi/config.vsh.yaml b/src/integrate/scvi/config.vsh.yaml
index b0adcc27443..f3cef372030 100644
--- a/src/integrate/scvi/config.vsh.yaml
+++ b/src/integrate/scvi/config.vsh.yaml
@@ -241,7 +241,7 @@ engines:
 
 runners:
 - type: executable
-  docker_run_args: ["--gpus all"]
+  # docker_run_args: ["--gpus all"]
 - type: nextflow
   directives:
     label: [midcpu, midmem, gpu, highdisk]

From 68c196060b764b50088e9b489c8a7c9e6ecdb9d5 Mon Sep 17 00:00:00 2001
From: Dries Schaumont <5946712+DriesSchaumont@users.noreply.github.com>
Date: Thu, 15 May 2025 07:12:06 +0000
Subject: [PATCH 3/3] Update CHANGELOG [ci skip]

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ece229b1d35..d9b97527d14 100644
--- a/CHANGELOG.md
+++ b/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