Merge branch 'main' into Ori-Dropping-Jax

Author: ori-kron-wis

CHANGELOG.md

@@ -13,6 +13,8 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
 
 #### Fixed
 
+- Fix list of metrics to be recorded in {class}`scvi.autotune.AutotuneExperiment`, {pr}`3816`.
+
 #### Changed
 
 #### Removed

pyproject.toml

@@ -59,8 +59,8 @@ tests = ["pytest", "pytest-pretty", "coverage", "scvi-tools[optional]"]
 editing = ["jupyter", "pre-commit"]
 dev = ["scvi-tools[editing,tests]"]
 test = ["scvi-tools[tests]"]
-cuda = ["torchvision", "torchaudio", "mlx[cuda]"]
-cuda13 = ["torchvision", "torchaudio", "mlx[cuda13]"]
+cuda = ["torchvision", "torchaudio"]
+cuda13 = ["torchvision", "torchaudio"]
 tpu = ["torch_xla[tpu]"]
 metal = ["torchvision", "torchaudio", "mlx-metal"]
 

src/scvi/autotune/_experiment.py

@@ -721,7 +721,7 @@ def metrics_callback(self) -> Callback:
         )
         on = "validation_end" if "validation" in self.metrics[0] else "train_end"
 
-        return callback_cls(metrics=self.metrics[0], on=on, save_checkpoints=self.save_checkpoints)
+        return callback_cls(metrics=self.metrics, on=on, save_checkpoints=self.save_checkpoints)
 
     @property
     def scib_metrics_callback(self) -> Callback:

tests/autotune/test_tune.py

@@ -45,6 +45,43 @@ def test_run_autotune_scvi_basic_adata(save_checkpoints: bool, metric: str, save
     assert isinstance(experiment.result_grid, ResultGrid)
 
 
+@pytest.mark.autotune
+def test_run_autotune_scvi_multiple_metrics(save_path: str):
+    from ray import tune
+    from ray.tune import ResultGrid
+
+    from scvi.autotune import AutotuneExperiment, run_autotune
+
+    settings.logging_dir = save_path
+    adata = synthetic_iid()
+    SCVI.setup_anndata(adata)
+
+    experiment = run_autotune(
+        SCVI,
+        adata,
+        metrics=["elbo_validation", "validation_loss"],
+        mode="min",
+        search_space={
+            "model_params": {
+                "n_hidden": tune.choice([1, 2]),
+            },
+            "train_params": {
+                "max_epochs": 1,
+            },
+        },
+        num_samples=2,
+        seed=0,
+        scheduler="asha",
+        searcher="hyperopt",
+        ignore_reinit_error=True,
+    )
+    assert isinstance(experiment, AutotuneExperiment)
+    assert isinstance(experiment.result_grid, ResultGrid)
+    result_df = experiment.result_grid.get_dataframe()
+    assert "elbo_validation" in result_df.columns
+    assert "validation_loss" in result_df.columns
+
+
 @pytest.mark.autotune
 @pytest.mark.parametrize("save_checkpoints", [True, False])
 def test_run_autotune_scvi_basic_mdata(save_checkpoints: bool, save_path: str):