Understanding the GPU related warning while setting up SCVI

Help scvi-tools
1

Hi,
I followed the exact steps documented in Installation - scvi-tools. I can see my gpu in nvidia-smi command. Yet whenever I am setting up the model as follows,

scvi.model.SCVI.setup_anndata(
    adata_liver,
    layer="counts",
    categorical_covariate_keys=["batch_id"]
)

I am meeting with the following error,

Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

I know these are warnings, but I would love to know what it really means or are there any downstream implications (e.g. not being able to utilize the GPU etc.)

Best,

2

These are related to our jax dependency. It will not be a problem that jax is not installed to use your gpu.