diff --git a/tuning/autotune_new_model.ipynb b/tuning/autotune_new_model.ipynb
index c9875b73..b1458d6f 100644
--- a/tuning/autotune_new_model.ipynb
+++ b/tuning/autotune_new_model.ipynb
@@ -42,33 +42,30 @@
                 "## Installing required packages"
             ]
         },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "Uncomment the following lines in Google Colab in order to install `scvi-tools`:"
+            ]
+        },
         {
             "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 1,
             "metadata": {},
             "outputs": [],
             "source": [
-                "!pip install --quiet hyperopt\n",
-                "!pip install --quiet \"ray[tune]\"\n",
-                "!pip install --quiet scvi-colab\n",
-                "from scvi_colab import install\n",
+                "# !pip install --quiet scvi-colab\n",
+                "# from scvi_colab import install\n",
                 "\n",
-                "install()"
+                "# install()"
             ]
         },
         {
             "cell_type": "code",
-            "execution_count": 1,
+            "execution_count": 2,
             "metadata": {},
-            "outputs": [
-                {
-                    "name": "stderr",
-                    "output_type": "stream",
-                    "text": [
-                        "Global seed set to 0\n"
-                    ]
-                }
-            ],
+            "outputs": [],
             "source": [
                 "import jax\n",
                 "import jax.numpy as jnp\n",
@@ -81,9 +78,17 @@
         },
         {
             "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 3,
             "metadata": {},
-            "outputs": [],
+            "outputs": [
+                {
+                    "name": "stdout",
+                    "output_type": "stream",
+                    "text": [
+                        "Last run with scvi-tools version: 1.0.3\n"
+                    ]
+                }
+            ],
             "source": [
                 "scvi.settings.seed = 0\n",
                 "print(\"Last run with scvi-tools version:\", scvi.__version__)"
@@ -107,7 +112,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 2,
+            "execution_count": 4,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -176,7 +181,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 3,
+            "execution_count": 5,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -217,15 +222,15 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 4,
+            "execution_count": 6,
             "metadata": {},
             "outputs": [
                 {
                     "name": "stderr",
                     "output_type": "stream",
                     "text": [
-                        "/home/martin/dev/scvi-tools/scvi/autotune/_manager.py:74: UserWarning: No default search space available for LassoTunable.\n",
-                        "  warnings.warn(\n"
+                        "/home/martin/dev/scvi-tools/scvi/autotune/_manager.py:60: UserWarning: No default search space available for LassoTunable.\n",
+                        "  self._defaults = self._get_defaults(self._model_cls)\n"
                     ]
                 },
                 {
@@ -334,7 +339,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 5,
+            "execution_count": 7,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -385,7 +390,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 6,
+            "execution_count": 8,
             "metadata": {},
             "outputs": [
                 {
@@ -496,7 +501,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 7,
+            "execution_count": 9,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -526,15 +531,15 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 8,
+            "execution_count": 10,
             "metadata": {},
             "outputs": [
                 {
                     "name": "stderr",
                     "output_type": "stream",
                     "text": [
-                        "/home/martin/dev/scvi-tools/scvi/autotune/_manager.py:74: UserWarning: No default search space available for LassoModel.\n",
-                        "  warnings.warn(\n"
+                        "/home/martin/dev/scvi-tools/scvi/autotune/_manager.py:60: UserWarning: No default search space available for LassoModel.\n",
+                        "  self._defaults = self._get_defaults(self._model_cls)\n"
                     ]
                 },
                 {
@@ -649,7 +654,7 @@
             "name": "python",
             "nbconvert_exporter": "python",
             "pygments_lexer": "ipython3",
-            "version": "3.10.8"
+            "version": "3.11.4"
         },
         "orig_nbformat": 4,
         "vscode": {
diff --git a/tuning/autotune_scvi.ipynb b/tuning/autotune_scvi.ipynb
index c53912ea..10c2f6eb 100644
--- a/tuning/autotune_scvi.ipynb
+++ b/tuning/autotune_scvi.ipynb
@@ -1,615 +1,661 @@
 {
-    "cells": [
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "# Model hyperparameter tuning with scVI"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "```{warning}\n",
-                "`scvi.autotune` development is still in progress. The API is subject to change.\n",
-                "```"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "Finding an effective set of model hyperparameters (e.g. learning rate, number of hidden layers, etc.) is an important component in training a model as its performance can be highly dependent on these non-trainable parameters. Manually tuning a model often involves picking a set of hyperparameters to search over and then evaluating different configurations over a validation set for a desired metric. This process can be time consuming and can require some prior intuition about a model and dataset pair, which is not always feasible.\n",
-                "\n",
-                "In this tutorial, we show how to use `scvi`'s [`autotune`](https://docs.scvi-tools.org/en/latest/api/user.html#model-hyperparameter-autotuning) module, which allows us to automatically find a good set of model hyperparameters using [Ray Tune](https://docs.ray.io/en/latest/tune/index.html). We will use `SCVI` and a subsample of the [heart cell atlas](https://www.heartcellatlas.org/#DataSources) for the task of batch correction, but the principles outlined here can be applied to any model and dataset. In particular, we will go through the following steps:\n",
-                "\n",
-                "1. Installing required packages\n",
-                "1. Loading and preprocessing the dataset\n",
-                "1. Defining the tuner and discovering hyperparameters\n",
-                "1. Running the tuner\n",
-                "1. Comparing latent spaces\n",
-                "1. Optional: Monitoring progress with Tensorboard\n",
-                "1. Optional: Tuning over integration metrics with `scib-metrics`"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "## Installing required packages"
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": null,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "!pip install --quiet hyperopt\n",
-                "!pip install --quiet \"ray[tune]\"\n",
-                "!pip install --quiet scvi-colab\n",
-                "from scvi_colab import install\n",
-                "\n",
-                "install()"
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 1,
-            "metadata": {},
-            "outputs": [
-                {
-                    "name": "stderr",
-                    "output_type": "stream",
-                    "text": [
-                        "Global seed set to 0\n"
-                    ]
-                }
-            ],
-            "source": [
-                "import ray\n",
-                "import scanpy as sc\n",
-                "import scvi\n",
-                "from ray import tune\n",
-                "from scvi import autotune"
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": null,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "scvi.settings.seed = 0\n",
-                "print(\"Last run with scvi-tools version:\", scvi.__version__)"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "## Loading and preprocessing the dataset"
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 2,
-            "metadata": {},
-            "outputs": [
-                {
-                    "name": "stdout",
-                    "output_type": "stream",
-                    "text": [
-                        "\u001b[34mINFO    \u001b[0m File data/hca_subsampled_20k.h5ad already downloaded                                                      \n"
-                    ]
-                },
-                {
-                    "data": {
-                        "text/plain": [
-                            "AnnData object with n_obs × n_vars = 18641 × 26662\n",
-                            "    obs: 'NRP', 'age_group', 'cell_source', 'cell_type', 'donor', 'gender', 'n_counts', 'n_genes', 'percent_mito', 'percent_ribo', 'region', 'sample', 'scrublet_score', 'source', 'type', 'version', 'cell_states', 'Used'\n",
-                            "    var: 'gene_ids-Harvard-Nuclei', 'feature_types-Harvard-Nuclei', 'gene_ids-Sanger-Nuclei', 'feature_types-Sanger-Nuclei', 'gene_ids-Sanger-Cells', 'feature_types-Sanger-Cells', 'gene_ids-Sanger-CD45', 'feature_types-Sanger-CD45', 'n_counts'\n",
-                            "    uns: 'cell_type_colors'"
-                        ]
-                    },
-                    "execution_count": 2,
-                    "metadata": {},
-                    "output_type": "execute_result"
-                }
-            ],
-            "source": [
-                "adata = scvi.data.heart_cell_atlas_subsampled()\n",
-                "adata"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "The only preprocessing step we will perform in this case will be to subsample the dataset for 2000 highly variable genes using `scanpy` for faster model training."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 3,
-            "metadata": {},
-            "outputs": [
-                {
-                    "data": {
-                        "text/plain": [
-                            "AnnData object with n_obs × n_vars = 18641 × 2000\n",
-                            "    obs: 'NRP', 'age_group', 'cell_source', 'cell_type', 'donor', 'gender', 'n_counts', 'n_genes', 'percent_mito', 'percent_ribo', 'region', 'sample', 'scrublet_score', 'source', 'type', 'version', 'cell_states', 'Used'\n",
-                            "    var: 'gene_ids-Harvard-Nuclei', 'feature_types-Harvard-Nuclei', 'gene_ids-Sanger-Nuclei', 'feature_types-Sanger-Nuclei', 'gene_ids-Sanger-Cells', 'feature_types-Sanger-Cells', 'gene_ids-Sanger-CD45', 'feature_types-Sanger-CD45', 'n_counts', 'highly_variable', 'highly_variable_rank', 'means', 'variances', 'variances_norm'\n",
-                            "    uns: 'cell_type_colors', 'hvg'"
-                        ]
-                    },
-                    "execution_count": 3,
-                    "metadata": {},
-                    "output_type": "execute_result"
-                }
-            ],
-            "source": [
-                "sc.pp.highly_variable_genes(adata, n_top_genes=2000, flavor=\"seurat_v3\", subset=True)\n",
-                "adata"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "## Defining the tuner and discovering hyperparameters"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "The first part of our workflow is the same as the standard `scvi-tools` workflow: we start with our desired model class, and we register our dataset with it using `setup_anndata`. All datasets must be registered prior to hyperparameter tuning."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 4,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "model_cls = scvi.model.SCVI\n",
-                "model_cls.setup_anndata(adata)"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "Our main entry point to the `autotune` module is the `ModelTuner` class, a wrapper around [`ray.tune.Tuner`](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tuner) with additional functionality specific to `scvi-tools`. We can define a new `ModelTuner` by providing it with our model class."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 5,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "scvi_tuner = autotune.ModelTuner(model_cls)"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "`ModelTuner` will register all tunable hyperparameters in `SCVI` -- these can be viewed by calling `info()`. By default, this method will display three tables:\n",
-                "\n",
-                "1. **Tunable hyperparameters**: The names of hyperparameters that can be tuned, their default values, and the internal classes they are defined in.\n",
-                "1. **Available metrics**: The metrics that can be used to evaluate the performance of the model. One of these must be provided when running the tuner.\n",
-                "1. **Default search space**: The default search space for the model class, which will be used if no search space is provided by the user."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 6,
-            "metadata": {},
-            "outputs": [
-                {
-                    "data": {
-                        "text/html": [
-                            "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">ModelTuner registry for SCVI\n",
-                            "</pre>\n"
-                        ],
-                        "text/plain": [
-                            "ModelTuner registry for SCVI\n"
-                        ]
-                    },
-                    "metadata": {},
-                    "output_type": "display_data"
-                },
-                {
-                    "data": {
-                        "text/html": [
-                            "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                  Tunable hyperparameters                  </span>\n",
-                            "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓\n",
-                            "┃<span style=\"font-weight: bold\">      Hyperparameter      </span>┃<span style=\"font-weight: bold\"> Default value </span>┃<span style=\"font-weight: bold\">    Source    </span>┃\n",
-                            "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_hidden         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      128      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_latent         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      10       </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_layers         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">       1       </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       dropout_rate       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.1      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">        dispersion        </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     gene      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">     gene_likelihood      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     zinb      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">   latent_distribution    </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">    normal     </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    encode_covariates     </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     False     </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> deeply_inject_covariates </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     True      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      use_batch_norm      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     both      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      use_layer_norm      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     none      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">        optimizer         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     Adam      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">            lr            </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     0.001     </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       weight_decay       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     1e-06     </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">           eps            </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     0.01      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    n_steps_kl_warmup     </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     None      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    n_epochs_kl_warmup    </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      400      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">   reduce_lr_on_plateau   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     False     </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">        lr_factor         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.6      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       lr_patience        </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      30       </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       lr_threshold       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.0      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">          lr_min          </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">       0       </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      max_kl_weight       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      1.0      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      min_kl_weight       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.0      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
-                            "└──────────────────────────┴───────────────┴──────────────┘\n",
-                            "</pre>\n"
-                        ],
-                        "text/plain": [
-                            "\u001b[3m                  Tunable hyperparameters                  \u001b[0m\n",
-                            "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓\n",
-                            "┃\u001b[1m \u001b[0m\u001b[1m     Hyperparameter     \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mDefault value\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m   Source   \u001b[0m\u001b[1m \u001b[0m┃\n",
-                            "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_hidden        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     128     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_latent        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     10      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_layers        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m      1      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      dropout_rate      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.1     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m       dispersion       \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    gene     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m    gene_likelihood     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    zinb     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m  latent_distribution   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m   normal    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   encode_covariates    \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    False    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mdeeply_inject_covariates\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    True     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     use_batch_norm     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    both     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     use_layer_norm     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    none     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m       optimizer        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    Adam     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m           lr           \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    0.001    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      weight_decay      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    1e-06    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m          eps           \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    0.01     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   n_steps_kl_warmup    \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    None     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   n_epochs_kl_warmup   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     400     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m  reduce_lr_on_plateau  \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    False    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m       lr_factor        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.6     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      lr_patience       \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     30      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      lr_threshold      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.0     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m         lr_min         \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m      0      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     max_kl_weight      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     1.0     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     min_kl_weight      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.0     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
-                            "└──────────────────────────┴───────────────┴──────────────┘\n"
-                        ]
-                    },
-                    "metadata": {},
-                    "output_type": "display_data"
-                },
-                {
-                    "data": {
-                        "text/html": [
-                            "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">       Available metrics        </span>\n",
-                            "┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓\n",
-                            "┃<span style=\"font-weight: bold\">     Metric      </span>┃<span style=\"font-weight: bold\">    Mode    </span>┃\n",
-                            "┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> validation_loss </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">    min     </span>│\n",
-                            "└─────────────────┴────────────┘\n",
-                            "</pre>\n"
-                        ],
-                        "text/plain": [
-                            "\u001b[3m       Available metrics        \u001b[0m\n",
-                            "┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓\n",
-                            "┃\u001b[1m \u001b[0m\u001b[1m    Metric     \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m   Mode   \u001b[0m\u001b[1m \u001b[0m┃\n",
-                            "┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mvalidation_loss\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m   min    \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
-                            "└─────────────────┴────────────┘\n"
-                        ]
-                    },
-                    "metadata": {},
-                    "output_type": "display_data"
-                },
-                {
-                    "data": {
-                        "text/html": [
-                            "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                         Default search space                         </span>\n",
-                            "┏━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓\n",
-                            "┃<span style=\"font-weight: bold\"> Hyperparameter </span>┃<span style=\"font-weight: bold\"> Sample function </span>┃<span style=\"font-weight: bold\">  Arguments  </span>┃<span style=\"font-weight: bold\"> Keyword arguments </span>┃\n",
-                            "┡━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩\n",
-                            "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    n_hidden    </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     choice      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> [[64, 128]] </span>│<span style=\"color: #ff8700; text-decoration-color: #ff8700\">        {}         </span>│\n",
-                            "└────────────────┴─────────────────┴─────────────┴───────────────────┘\n",
-                            "</pre>\n"
-                        ],
-                        "text/plain": [
-                            "\u001b[3m                         Default search space                         \u001b[0m\n",
-                            "┏━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓\n",
-                            "┃\u001b[1m \u001b[0m\u001b[1mHyperparameter\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mSample function\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m Arguments \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mKeyword arguments\u001b[0m\u001b[1m \u001b[0m┃\n",
-                            "┡━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩\n",
-                            "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   n_hidden   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    choice     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m[[64, 128]]\u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;208m \u001b[0m\u001b[38;5;208m       {}        \u001b[0m\u001b[38;5;208m \u001b[0m│\n",
-                            "└────────────────┴─────────────────┴─────────────┴───────────────────┘\n"
-                        ]
-                    },
-                    "metadata": {},
-                    "output_type": "display_data"
-                }
-            ],
-            "source": [
-                "scvi_tuner.info()"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "## Running the tuner"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "Now that we know what hyperparameters are available to us, we can define a search space using the [search space API](https://docs.ray.io/en/latest/tune/api/search_space.html) in `ray.tune`. For this tutorial, we choose a simple search space with two model hyperparameters and one training plan hyperparameter. These can all be combined into a single dictionary that we pass into the `fit` method."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 7,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "search_space = {\n",
-                "    \"n_hidden\": tune.choice([64, 128, 256]),\n",
-                "    \"n_layers\": tune.choice([1, 2, 3]),\n",
-                "    \"lr\": tune.loguniform(1e-4, 1e-2),\n",
-                "}"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "There are a couple more arguments we should be aware of before fitting the tuner:\n",
-                "\n",
-                "- `num_samples`: The total number of hyperparameter sets to sample from `search_space`. This is the total number of models that will be trained.\n",
-                "\n",
-                "  For example, if we set `num_samples=2`, we might sample two models with the following hyperparameter configurations:\n",
-                "\n",
-                "  ```python\n",
-                "  model1 = {\n",
-                "      \"n_hidden\": 64,\n",
-                "      \"n_layers\": 1,\n",
-                "      \"lr\": 0.001,\n",
-                "  }\n",
-                "  model2 = {\n",
-                "      \"n_hidden\": 64,\n",
-                "      \"n_layers\": 3,\n",
-                "      \"lr\": 0.0001,\n",
-                "  }\n",
-                "  ```\n",
-                "\n",
-                "- `max_epochs`: The maximum number of epochs to train each model for.\n",
-                "\n",
-                "  Note: This does not mean that each model will be trained for `max_epochs`. Depending on the scheduler used, some trials are likely to be stopped early.\n",
-                "\n",
-                "- `resources`: A dictionary of maximum resources to allocate for the whole experiment. This allows us to run concurrent trials on limited hardware.\n",
-                "\n",
-                "Now, we can call `fit` on the tuner to start the hyperparameter sweep. This will return a `TuneAnalysis` dataclass, which will contain the best set of hyperparameters, as well as other information."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 8,
-            "metadata": {},
-            "outputs": [
-                {
-                    "data": {
-                        "text/html": [
-                            "<div class=\"tuneStatus\">\n",
-                            "  <div style=\"display: flex;flex-direction: row\">\n",
-                            "    <div style=\"display: flex;flex-direction: column;\">\n",
-                            "      <h3>Tune Status</h3>\n",
-                            "      <table>\n",
-                            "<tbody>\n",
-                            "<tr><td>Current time:</td><td>2023-01-30 13:44:04</td></tr>\n",
-                            "<tr><td>Running for: </td><td>00:03:01.04        </td></tr>\n",
-                            "<tr><td>Memory:      </td><td>10.9/125.7 GiB     </td></tr>\n",
-                            "</tbody>\n",
-                            "</table>\n",
-                            "    </div>\n",
-                            "    <div class=\"vDivider\"></div>\n",
-                            "    <div class=\"systemInfo\">\n",
-                            "      <h3>System Info</h3>\n",
-                            "      Using AsyncHyperBand: num_stopped=5<br>Bracket: Iter 64.000: -466.8729553222656 | Iter 32.000: -467.76768493652344 | Iter 16.000: -473.09934997558594 | Iter 8.000: -481.9038391113281 | Iter 4.000: -493.9166717529297 | Iter 2.000: -516.9609985351562 | Iter 1.000: -559.62353515625<br>Resources requested: 0/20 CPUs, 0/1 GPUs, 0.0/74.12 GiB heap, 0.0/35.76 GiB objects (0.0/1.0 accelerator_type:G)\n",
-                            "    </div>\n",
-                            "    \n",
-                            "  </div>\n",
-                            "  <div class=\"hDivider\"></div>\n",
-                            "  <div class=\"trialStatus\">\n",
-                            "    <h3>Trial Status</h3>\n",
-                            "    <table>\n",
-                            "<thead>\n",
-                            "<tr><th>Trial name            </th><th>status    </th><th>loc                  </th><th style=\"text-align: right;\">  n_hidden</th><th style=\"text-align: right;\">  n_layers</th><th style=\"text-align: right;\">         lr</th><th style=\"text-align: right;\">  validation_loss</th></tr>\n",
-                            "</thead>\n",
-                            "<tbody>\n",
-                            "<tr><td>_trainable_cb47e_00000</td><td>TERMINATED</td><td>128.32.142.133:515229</td><td style=\"text-align: right;\">       128</td><td style=\"text-align: right;\">         1</td><td style=\"text-align: right;\">0.000109758</td><td style=\"text-align: right;\">          468.995</td></tr>\n",
-                            "<tr><td>_trainable_cb47e_00001</td><td>TERMINATED</td><td>128.32.142.133:515229</td><td style=\"text-align: right;\">       256</td><td style=\"text-align: right;\">         1</td><td style=\"text-align: right;\">0.00906226 </td><td style=\"text-align: right;\">          469.203</td></tr>\n",
-                            "<tr><td>_trainable_cb47e_00002</td><td>TERMINATED</td><td>128.32.142.133:515229</td><td style=\"text-align: right;\">       128</td><td style=\"text-align: right;\">         3</td><td style=\"text-align: right;\">0.00276226 </td><td style=\"text-align: right;\">          616.49 </td></tr>\n",
-                            "<tr><td>_trainable_cb47e_00003</td><td>TERMINATED</td><td>128.32.142.133:515229</td><td style=\"text-align: right;\">       128</td><td style=\"text-align: right;\">         1</td><td style=\"text-align: right;\">0.00110585 </td><td style=\"text-align: right;\">          516.961</td></tr>\n",
-                            "<tr><td>_trainable_cb47e_00004</td><td>TERMINATED</td><td>128.32.142.133:515229</td><td style=\"text-align: right;\">       256</td><td style=\"text-align: right;\">         3</td><td style=\"text-align: right;\">0.000817148</td><td style=\"text-align: right;\">          595.537</td></tr>\n",
-                            "</tbody>\n",
-                            "</table>\n",
-                            "  </div>\n",
-                            "</div>\n",
-                            "<style>\n",
-                            ".tuneStatus {\n",
-                            "  color: var(--jp-ui-font-color1);\n",
-                            "}\n",
-                            ".tuneStatus .systemInfo {\n",
-                            "  display: flex;\n",
-                            "  flex-direction: column;\n",
-                            "}\n",
-                            ".tuneStatus td {\n",
-                            "  white-space: nowrap;\n",
-                            "}\n",
-                            ".tuneStatus .trialStatus {\n",
-                            "  display: flex;\n",
-                            "  flex-direction: column;\n",
-                            "}\n",
-                            ".tuneStatus h3 {\n",
-                            "  font-weight: bold;\n",
-                            "}\n",
-                            ".tuneStatus .hDivider {\n",
-                            "  border-bottom-width: var(--jp-border-width);\n",
-                            "  border-bottom-color: var(--jp-border-color0);\n",
-                            "  border-bottom-style: solid;\n",
-                            "}\n",
-                            ".tuneStatus .vDivider {\n",
-                            "  border-left-width: var(--jp-border-width);\n",
-                            "  border-left-color: var(--jp-border-color0);\n",
-                            "  border-left-style: solid;\n",
-                            "  margin: 0.5em 1em 0.5em 1em;\n",
-                            "}\n",
-                            "</style>\n"
-                        ],
-                        "text/plain": [
-                            "<IPython.core.display.HTML object>"
-                        ]
-                    },
-                    "metadata": {},
-                    "output_type": "display_data"
-                },
-                {
-                    "name": "stderr",
-                    "output_type": "stream",
-                    "text": [
-                        "2023-01-30 13:44:04,458\tINFO tune.py:762 -- Total run time: 181.26 seconds (181.03 seconds for the tuning loop).\n"
-                    ]
-                }
-            ],
-            "source": [
-                "ray.init(log_to_driver=False)\n",
-                "results = scvi_tuner.fit(\n",
-                "    adata,\n",
-                "    metric=\"validation_loss\",\n",
-                "    search_space=search_space,\n",
-                "    num_samples=5,\n",
-                "    max_epochs=100,\n",
-                "    resources={\"cpu\": 20, \"gpu\": 1},\n",
-                ")"
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 9,
-            "metadata": {},
-            "outputs": [
-                {
-                    "name": "stdout",
-                    "output_type": "stream",
-                    "text": [
-                        "{'n_hidden': 128, 'n_layers': 1}\n",
-                        "{'plan_kwargs': {}}\n"
-                    ]
-                }
-            ],
-            "source": [
-                "print(results.model_kwargs)\n",
-                "print(results.train_kwargs)"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "## Comparing latent spaces"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "Work in progress: please check back in the next release!"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "## Optional: Monitoring progress with Tensorboard"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "Work in progress: please check back in the next release!"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "## Optional: Tuning over integration metrics with `scib-metrics`"
-            ]
-        },
-        {
-            "attachments": {},
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "Work in progress: please check back in the next release!"
-            ]
-        }
-    ],
-    "metadata": {
-        "kernelspec": {
-            "display_name": "scvi-gpu",
-            "language": "python",
-            "name": "python3"
-        },
-        "language_info": {
-            "codemirror_mode": {
-                "name": "ipython",
-                "version": 3
-            },
-            "file_extension": ".py",
-            "mimetype": "text/x-python",
-            "name": "python",
-            "nbconvert_exporter": "python",
-            "pygments_lexer": "ipython3",
-            "version": "3.10.8"
-        },
-        "orig_nbformat": 4,
-        "vscode": {
-            "interpreter": {
-                "hash": "2f978838050607ec9770689d8200902a4128a2ce208b502e911dd714d57e924e"
-            }
-        }
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Model hyperparameter tuning with scVI"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```{warning}\n",
+    "`scvi.autotune` development is still in progress. The API is subject to change.\n",
+    "```"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Finding an effective set of model hyperparameters (e.g. learning rate, number of hidden layers, etc.) is an important component in training a model as its performance can be highly dependent on these non-trainable parameters. Manually tuning a model often involves picking a set of hyperparameters to search over and then evaluating different configurations over a validation set for a desired metric. This process can be time consuming and can require some prior intuition about a model and dataset pair, which is not always feasible.\n",
+    "\n",
+    "In this tutorial, we show how to use `scvi`'s [`autotune`](https://docs.scvi-tools.org/en/latest/api/user.html#model-hyperparameter-autotuning) module, which allows us to automatically find a good set of model hyperparameters using [Ray Tune](https://docs.ray.io/en/latest/tune/index.html). We will use `SCVI` and a subsample of the [heart cell atlas](https://www.heartcellatlas.org/#DataSources) for the task of batch correction, but the principles outlined here can be applied to any model and dataset. In particular, we will go through the following steps:\n",
+    "\n",
+    "1. Installing required packages\n",
+    "1. Loading and preprocessing the dataset\n",
+    "1. Defining the tuner and discovering hyperparameters\n",
+    "1. Running the tuner\n",
+    "1. Comparing latent spaces\n",
+    "1. Optional: Monitoring progress with Tensorboard\n",
+    "1. Optional: Tuning over integration metrics with `scib-metrics`"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Installing required packages"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Uncomment the following lines in Google Colab in order to install `scvi-tools`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# !pip install --quiet scvi-colab\n",
+    "# from scvi_colab import install\n",
+    "\n",
+    "# install()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tempfile\n",
+    "\n",
+    "import ray\n",
+    "import scanpy as sc\n",
+    "import scvi\n",
+    "import torch\n",
+    "from ray import tune\n",
+    "from scvi import autotune"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Global seed set to 0\n"
+     ]
     },
-    "nbformat": 4,
-    "nbformat_minor": 2
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Last run with scvi-tools version: 1.0.3\n"
+     ]
+    }
+   ],
+   "source": [
+    "scvi.settings.seed = 0\n",
+    "print(\"Last run with scvi-tools version:\", scvi.__version__)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can modify `save_dir` below to change where the data files for this tutorial are saved."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sc.set_figure_params(figsize=(4, 4), frameon=False)\n",
+    "torch.set_float32_matmul_precision(\"high\")\n",
+    "save_dir = tempfile.TemporaryDirectory()\n",
+    "\n",
+    "%config InlineBackend.print_figure_kwargs={\"facecolor\" : \"w\"}\n",
+    "%config InlineBackend.figure_format=\"retina\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Loading and preprocessing the dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34mINFO    \u001b[0m Downloading file at \u001b[35m/tmp/tmp51dvndw4/\u001b[0m\u001b[95mhca_subsampled_20k.h5ad\u001b[0m                                              \n",
+      "Downloading...: 100%|██████████| 65714/65714.0 [00:00<00:00, 110628.59it/s]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "AnnData object with n_obs × n_vars = 18641 × 26662\n",
+       "    obs: 'NRP', 'age_group', 'cell_source', 'cell_type', 'donor', 'gender', 'n_counts', 'n_genes', 'percent_mito', 'percent_ribo', 'region', 'sample', 'scrublet_score', 'source', 'type', 'version', 'cell_states', 'Used'\n",
+       "    var: 'gene_ids-Harvard-Nuclei', 'feature_types-Harvard-Nuclei', 'gene_ids-Sanger-Nuclei', 'feature_types-Sanger-Nuclei', 'gene_ids-Sanger-Cells', 'feature_types-Sanger-Cells', 'gene_ids-Sanger-CD45', 'feature_types-Sanger-CD45', 'n_counts'\n",
+       "    uns: 'cell_type_colors'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "adata = scvi.data.heart_cell_atlas_subsampled(save_path=save_dir.name)\n",
+    "adata"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The only preprocessing step we will perform in this case will be to subsample the dataset for 2000 highly variable genes using `scanpy` for faster model training."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "AnnData object with n_obs × n_vars = 18641 × 2000\n",
+       "    obs: 'NRP', 'age_group', 'cell_source', 'cell_type', 'donor', 'gender', 'n_counts', 'n_genes', 'percent_mito', 'percent_ribo', 'region', 'sample', 'scrublet_score', 'source', 'type', 'version', 'cell_states', 'Used'\n",
+       "    var: 'gene_ids-Harvard-Nuclei', 'feature_types-Harvard-Nuclei', 'gene_ids-Sanger-Nuclei', 'feature_types-Sanger-Nuclei', 'gene_ids-Sanger-Cells', 'feature_types-Sanger-Cells', 'gene_ids-Sanger-CD45', 'feature_types-Sanger-CD45', 'n_counts', 'highly_variable', 'highly_variable_rank', 'means', 'variances', 'variances_norm'\n",
+       "    uns: 'cell_type_colors', 'hvg'"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "sc.pp.highly_variable_genes(adata, n_top_genes=2000, flavor=\"seurat_v3\", subset=True)\n",
+    "adata"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Defining the tuner and discovering hyperparameters"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The first part of our workflow is the same as the standard `scvi-tools` workflow: we start with our desired model class, and we register our dataset with it using `setup_anndata`. All datasets must be registered prior to hyperparameter tuning."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_cls = scvi.model.SCVI\n",
+    "model_cls.setup_anndata(adata)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Our main entry point to the `autotune` module is the `ModelTuner` class, a wrapper around [`ray.tune.Tuner`](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tuner) with additional functionality specific to `scvi-tools`. We can define a new `ModelTuner` by providing it with our model class."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "scvi_tuner = autotune.ModelTuner(model_cls)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "`ModelTuner` will register all tunable hyperparameters in `SCVI` -- these can be viewed by calling `info()`. By default, this method will display three tables:\n",
+    "\n",
+    "1. **Tunable hyperparameters**: The names of hyperparameters that can be tuned, their default values, and the internal classes they are defined in.\n",
+    "1. **Available metrics**: The metrics that can be used to evaluate the performance of the model. One of these must be provided when running the tuner.\n",
+    "1. **Default search space**: The default search space for the model class, which will be used if no search space is provided by the user."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">ModelTuner registry for SCVI\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "ModelTuner registry for SCVI\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                  Tunable hyperparameters                  </span>\n",
+       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓\n",
+       "┃<span style=\"font-weight: bold\">      Hyperparameter      </span>┃<span style=\"font-weight: bold\"> Default value </span>┃<span style=\"font-weight: bold\">    Source    </span>┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_hidden         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      128      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_latent         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      10       </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_layers         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">       1       </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       dropout_rate       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.1      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">        dispersion        </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     gene      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">     log_variational      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     True      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">     gene_likelihood      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     zinb      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">   latent_distribution    </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">    normal     </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    encode_covariates     </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     False     </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> deeply_inject_covariates </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     True      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      use_batch_norm      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     both      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      use_layer_norm      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     none      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">  use_observed_lib_size   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     True      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      var_activation      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     None      </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     VAE      </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">        optimizer         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     Adam      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">            lr            </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     0.001     </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       weight_decay       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     1e-06     </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">           eps            </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     0.01      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    n_steps_kl_warmup     </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     None      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    n_epochs_kl_warmup    </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      400      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">   reduce_lr_on_plateau   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     False     </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">        lr_factor         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.6      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       lr_patience        </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      30       </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       lr_threshold       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.0      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">          lr_min          </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">       0       </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      max_kl_weight       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      1.0      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      min_kl_weight       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">      0.0      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> TrainingPlan </span>│\n",
+       "└──────────────────────────┴───────────────┴──────────────┘\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[3m                  Tunable hyperparameters                  \u001b[0m\n",
+       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓\n",
+       "┃\u001b[1m \u001b[0m\u001b[1m     Hyperparameter     \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mDefault value\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m   Source   \u001b[0m\u001b[1m \u001b[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_hidden        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     128     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_latent        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     10      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_layers        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m      1      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      dropout_rate      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.1     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m       dispersion       \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    gene     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m    log_variational     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    True     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m    gene_likelihood     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    zinb     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m  latent_distribution   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m   normal    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   encode_covariates    \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    False    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mdeeply_inject_covariates\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    True     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     use_batch_norm     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    both     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     use_layer_norm     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    none     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m use_observed_lib_size  \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    True     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     var_activation     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    None     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    VAE     \u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m       optimizer        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    Adam     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m           lr           \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    0.001    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      weight_decay      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    1e-06    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m          eps           \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    0.01     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   n_steps_kl_warmup    \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    None     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   n_epochs_kl_warmup   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     400     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m  reduce_lr_on_plateau  \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    False    \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m       lr_factor        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.6     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      lr_patience       \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     30      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      lr_threshold      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.0     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m         lr_min         \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m      0      \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     max_kl_weight      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     1.0     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     min_kl_weight      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m     0.0     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32mTrainingPlan\u001b[0m\u001b[32m \u001b[0m│\n",
+       "└──────────────────────────┴───────────────┴──────────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">       Available metrics        </span>\n",
+       "┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓\n",
+       "┃<span style=\"font-weight: bold\">     Metric      </span>┃<span style=\"font-weight: bold\">    Mode    </span>┃\n",
+       "┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> validation_loss </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">    min     </span>│\n",
+       "└─────────────────┴────────────┘\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[3m       Available metrics        \u001b[0m\n",
+       "┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓\n",
+       "┃\u001b[1m \u001b[0m\u001b[1m    Metric     \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m   Mode   \u001b[0m\u001b[1m \u001b[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mvalidation_loss\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m   min    \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
+       "└─────────────────┴────────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                         Default search space                         </span>\n",
+       "┏━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃<span style=\"font-weight: bold\"> Hyperparameter </span>┃<span style=\"font-weight: bold\"> Sample function </span>┃<span style=\"font-weight: bold\">  Arguments  </span>┃<span style=\"font-weight: bold\"> Keyword arguments </span>┃\n",
+       "┡━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩\n",
+       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    n_hidden    </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">     choice      </span>│<span style=\"color: #008000; text-decoration-color: #008000\"> [[64, 128]] </span>│<span style=\"color: #ff8700; text-decoration-color: #ff8700\">        {}         </span>│\n",
+       "└────────────────┴─────────────────┴─────────────┴───────────────────┘\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[3m                         Default search space                         \u001b[0m\n",
+       "┏━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃\u001b[1m \u001b[0m\u001b[1mHyperparameter\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mSample function\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m Arguments \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mKeyword arguments\u001b[0m\u001b[1m \u001b[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩\n",
+       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   n_hidden   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m    choice     \u001b[0m\u001b[38;5;128m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m[[64, 128]]\u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;208m \u001b[0m\u001b[38;5;208m       {}        \u001b[0m\u001b[38;5;208m \u001b[0m│\n",
+       "└────────────────┴─────────────────┴─────────────┴───────────────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "scvi_tuner.info()"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Running the tuner"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now that we know what hyperparameters are available to us, we can define a search space using the [search space API](https://docs.ray.io/en/latest/tune/api/search_space.html) in `ray.tune`. For this tutorial, we choose a simple search space with two model hyperparameters and one training plan hyperparameter. These can all be combined into a single dictionary that we pass into the `fit` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "search_space = {\n",
+    "    \"n_hidden\": tune.choice([64, 128, 256]),\n",
+    "    \"n_layers\": tune.choice([1, 2, 3]),\n",
+    "    \"lr\": tune.loguniform(1e-4, 1e-2),\n",
+    "}"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "There are a couple more arguments we should be aware of before fitting the tuner:\n",
+    "\n",
+    "- `num_samples`: The total number of hyperparameter sets to sample from `search_space`. This is the total number of models that will be trained.\n",
+    "\n",
+    "  For example, if we set `num_samples=2`, we might sample two models with the following hyperparameter configurations:\n",
+    "\n",
+    "  ```python\n",
+    "  model1 = {\n",
+    "      \"n_hidden\": 64,\n",
+    "      \"n_layers\": 1,\n",
+    "      \"lr\": 0.001,\n",
+    "  }\n",
+    "  model2 = {\n",
+    "      \"n_hidden\": 64,\n",
+    "      \"n_layers\": 3,\n",
+    "      \"lr\": 0.0001,\n",
+    "  }\n",
+    "  ```\n",
+    "\n",
+    "- `max_epochs`: The maximum number of epochs to train each model for.\n",
+    "\n",
+    "  Note: This does not mean that each model will be trained for `max_epochs`. Depending on the scheduler used, some trials are likely to be stopped early.\n",
+    "\n",
+    "- `resources`: A dictionary of maximum resources to allocate for the whole experiment. This allows us to run concurrent trials on limited hardware.\n",
+    "\n",
+    "Now, we can call `fit` on the tuner to start the hyperparameter sweep. This will return a `TuneAnalysis` dataclass, which will contain the best set of hyperparameters, as well as other information."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div class=\"tuneStatus\">\n",
+       "  <div style=\"display: flex;flex-direction: row\">\n",
+       "    <div style=\"display: flex;flex-direction: column;\">\n",
+       "      <h3>Tune Status</h3>\n",
+       "      <table>\n",
+       "<tbody>\n",
+       "<tr><td>Current time:</td><td>2023-07-27 03:16:53</td></tr>\n",
+       "<tr><td>Running for: </td><td>00:02:40.69        </td></tr>\n",
+       "<tr><td>Memory:      </td><td>34.8/125.7 GiB     </td></tr>\n",
+       "</tbody>\n",
+       "</table>\n",
+       "    </div>\n",
+       "    <div class=\"vDivider\"></div>\n",
+       "    <div class=\"systemInfo\">\n",
+       "      <h3>System Info</h3>\n",
+       "      Using AsyncHyperBand: num_stopped=5<br>Bracket: Iter 64.000: -451.93446350097656 | Iter 32.000: -453.1311798095703 | Iter 16.000: -459.6698303222656 | Iter 8.000: -469.3856506347656 | Iter 4.000: -481.6686248779297 | Iter 2.000: -517.1959228515625 | Iter 1.000: -642.4600219726562<br>Logical resource usage: 10.0/20 CPUs, 1.0/1 GPUs (0.0/1.0 accelerator_type:G)\n",
+       "    </div>\n",
+       "    \n",
+       "  </div>\n",
+       "  <div class=\"hDivider\"></div>\n",
+       "  <div class=\"trialStatus\">\n",
+       "    <h3>Trial Status</h3>\n",
+       "    <table>\n",
+       "<thead>\n",
+       "<tr><th>Trial name         </th><th>status    </th><th>loc                  </th><th style=\"text-align: right;\">  n_hidden</th><th style=\"text-align: right;\">  n_layers</th><th style=\"text-align: right;\">         lr</th><th style=\"text-align: right;\">  validation_loss</th></tr>\n",
+       "</thead>\n",
+       "<tbody>\n",
+       "<tr><td>_trainable_9ca2fa7c</td><td>TERMINATED</td><td>128.32.142.133:965309</td><td style=\"text-align: right;\">       256</td><td style=\"text-align: right;\">         3</td><td style=\"text-align: right;\">0.000589003</td><td style=\"text-align: right;\">          455.212</td></tr>\n",
+       "<tr><td>_trainable_7d98a8b2</td><td>TERMINATED</td><td>128.32.142.133:965309</td><td style=\"text-align: right;\">        64</td><td style=\"text-align: right;\">         2</td><td style=\"text-align: right;\">0.000350049</td><td style=\"text-align: right;\">          737.187</td></tr>\n",
+       "<tr><td>_trainable_ff1772f6</td><td>TERMINATED</td><td>128.32.142.133:965309</td><td style=\"text-align: right;\">        64</td><td style=\"text-align: right;\">         3</td><td style=\"text-align: right;\">0.000918661</td><td style=\"text-align: right;\">          556.365</td></tr>\n",
+       "<tr><td>_trainable_5fa3a6f6</td><td>TERMINATED</td><td>128.32.142.133:965309</td><td style=\"text-align: right;\">       256</td><td style=\"text-align: right;\">         3</td><td style=\"text-align: right;\">0.000161977</td><td style=\"text-align: right;\">          770.872</td></tr>\n",
+       "<tr><td>_trainable_53fa3aa1</td><td>TERMINATED</td><td>128.32.142.133:965309</td><td style=\"text-align: right;\">       128</td><td style=\"text-align: right;\">         1</td><td style=\"text-align: right;\">0.00122786 </td><td style=\"text-align: right;\">          452.142</td></tr>\n",
+       "</tbody>\n",
+       "</table>\n",
+       "  </div>\n",
+       "</div>\n",
+       "<style>\n",
+       ".tuneStatus {\n",
+       "  color: var(--jp-ui-font-color1);\n",
+       "}\n",
+       ".tuneStatus .systemInfo {\n",
+       "  display: flex;\n",
+       "  flex-direction: column;\n",
+       "}\n",
+       ".tuneStatus td {\n",
+       "  white-space: nowrap;\n",
+       "}\n",
+       ".tuneStatus .trialStatus {\n",
+       "  display: flex;\n",
+       "  flex-direction: column;\n",
+       "}\n",
+       ".tuneStatus h3 {\n",
+       "  font-weight: bold;\n",
+       "}\n",
+       ".tuneStatus .hDivider {\n",
+       "  border-bottom-width: var(--jp-border-width);\n",
+       "  border-bottom-color: var(--jp-border-color0);\n",
+       "  border-bottom-style: solid;\n",
+       "}\n",
+       ".tuneStatus .vDivider {\n",
+       "  border-left-width: var(--jp-border-width);\n",
+       "  border-left-color: var(--jp-border-color0);\n",
+       "  border-left-style: solid;\n",
+       "  margin: 0.5em 1em 0.5em 1em;\n",
+       "}\n",
+       "</style>\n"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/martin/bin/mambaforge/envs/scvi-tools-dev/lib/python3.11/site-packages/ray/air/config.py:803: UserWarning: Setting a `RunConfig.local_dir` is deprecated and will be removed in the future. If you are not using remote storage,set the `RunConfig.storage_path` instead. Otherwise, set the`RAY_AIR_LOCAL_CACHE_DIR` environment variable to control the local cache location.\n",
+      "  warnings.warn(\n",
+      "2023-07-27 03:14:12,825\tINFO tune.py:657 -- [output] This uses the legacy output and progress reporter, as Jupyter notebooks are not supported by the new engine, yet. For more information, please see https://github.com/ray-project/ray/issues/36949\n",
+      "2023-07-27 03:16:53,535\tINFO tune.py:1148 -- Total run time: 160.71 seconds (160.69 seconds for the tuning loop).\n"
+     ]
+    }
+   ],
+   "source": [
+    "ray.init(log_to_driver=False)\n",
+    "results = scvi_tuner.fit(\n",
+    "    adata,\n",
+    "    metric=\"validation_loss\",\n",
+    "    search_space=search_space,\n",
+    "    num_samples=5,\n",
+    "    max_epochs=100,\n",
+    "    resources={\"cpu\": 10, \"gpu\": 1},\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'n_hidden': 128, 'n_layers': 1}\n",
+      "{'plan_kwargs': {'lr': 0.001227857394187722}}\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(results.model_kwargs)\n",
+    "print(results.train_kwargs)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Comparing latent spaces"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Work in progress: please check back in the next release!"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Optional: Monitoring progress with Tensorboard"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Work in progress: please check back in the next release!"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Optional: Tuning over integration metrics with `scib-metrics`"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Work in progress: please check back in the next release!"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "scvi-gpu",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "2f978838050607ec9770689d8200902a4128a2ce208b502e911dd714d57e924e"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
 }