Add cuda accelerated jax for scvi-tools#1028

dorien-er merged 3 commits into

mainopenpipelines-bio/openpipeline:mainfrom

fix_missing_cuda_jaxopenpipelines-bio/openpipeline:fix_missing_cuda_jaxCopy head branch name to clipboard