docs: updated for mlx (#3766)

ori-kron-wis · web-flow · commit 54fd6135263b · 2026-04-03T17:37:16.000+03:00
diff --git a/docs/conf.py b/docs/conf.py
@@ -93,7 +93,7 @@
 typehints_defaults = "braces"
 autodoc_mock_imports = []
 if os.environ.get("READTHEDOCS") == "True":
-    autodoc_mock_imports += ["hyperopt", "ray", "ray.tune", "scib_metrics", "muon"]
+    autodoc_mock_imports += ["hyperopt", "ray", "ray.tune", "scib_metrics", "muon", "mlx"]
     try:
         import scvi.autotune
     except Exception:
diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md
@@ -189,7 +189,7 @@ scvi-tools is composed of models that can perform one or many analysis tasks. In
    * - :doc:`/user_guide/models/amortizedlda`
      - Topic modeling
      - :cite:p:`Blei03`
-   * - :doc:`/tutorials/notebooks/hub/scvi_hub_intro_and_download`
+   * - :doc:`Scvi-Hub </tutorials/notebooks/hub/scvi_hub_intro_and_download>`
      - Scvi-hub: an actionable repository for model-driven single-cell analysis usign Hugging Face Hub
      - :cite:p:`Ergen25-2`
 
diff --git a/docs/user_guide/models/resolvi.md b/docs/user_guide/models/resolvi.md
@@ -84,7 +84,7 @@ The latent variables, along with their description, are summarized in the follow
 
 ```{eval-rst}
 .. list-table::
-   :widths: 20 90 15of ce
+   :widths: 20 90 15 45
    :header-rows: 1
 
    * - Latent variable
diff --git a/src/scvi/external/contrastivevi/_module.py b/src/scvi/external/contrastivevi/_module.py
@@ -545,6 +545,7 @@ def loss(
         Returns
         -------
         An scvi.module.base.LossOutput instance that records the following:
+
         loss
             One-dimensional tensor for overall loss used for optimization.
         reconstruction_loss
diff --git a/src/scvi/external/cytovi/_model.py b/src/scvi/external/cytovi/_model.py
@@ -102,7 +102,7 @@ class CYTOVI(
         required when analyzing overlapping panels with missing values.
     encoder_marker_list
         Optional list of markers to use for encoding. Must be a subset of backbone markers
-          if `encode_backbone_only` is True.
+        if `encode_backbone_only` is True.
     prior_mixture
         If True, uses a mixture of Gaussians as a prior in the latent space (MoG prior).
     prior_mixture_k
@@ -126,7 +126,7 @@ class CYTOVI(
     An adversarial classifier loss can be used to encourage batch-invariance in the latent space.
     If the data includes missing values, ensure that `nan_layer` is correctly registered using
     :meth:`~scvi.external.CYTOVI.setup_anndata`. This is handled automatically when using
-     scvi.external.cytovi.merge_batches().
+    scvi.external.cytovi.merge_batches().
 
     See further usage examples in the following tutorials:
 
diff --git a/src/scvi/external/decipher/_model.py b/src/scvi/external/decipher/_model.py
@@ -292,6 +292,7 @@ def compute_gene_patterns(
         -------
         The gene patterns for the trajectory.
         Dictionary keys:
+
             - `mean`: the mean gene expression pattern
             - `q25`: the 25% quantile of the gene expression pattern
             - `q75`: the 75% quantile of the gene expression pattern
diff --git a/src/scvi/external/methylvi/_base_components.py b/src/scvi/external/methylvi/_base_components.py
@@ -308,6 +308,7 @@ def differential_methylation(
         Returns
         -------
         Differential methylation DataFrame with the following columns:
+
         proba_de
             the probability of the region being differentially methylated
         is_de_fdr
diff --git a/src/scvi/external/methylvi/_methylanvi_model.py b/src/scvi/external/methylvi/_methylanvi_model.py
@@ -202,18 +202,16 @@ def setup_mudata(
 
         Examples
         --------
-        METHYLANVI.setup_mudata(
-            mdata,
-            mc_layer="mc",
-            cov_layer="cov",
-            labels_key="CellType",
-            unlabeled_category="Unknown",
-            methylation_contexts=["mCG", "mCH"],
-            categorical_covariate_keys=["Platform"],
-            modalities={
-                "categorical_covariate_keys": "mCG"
-            },
-        )
+        >>> METHYLANVI.setup_mudata(
+        ...     mdata,
+        ...     mc_layer="mc",
+        ...     cov_layer="cov",
+        ...     labels_key="CellType",
+        ...     unlabeled_category="Unknown",
+        ...     methylation_contexts=["mCG", "mCH"],
+        ...     categorical_covariate_keys=["Platform"],
+        ...     modalities={"categorical_covariate_keys": "mCG"},
+        ... )
 
         """
         if modalities is None:
diff --git a/src/scvi/external/methylvi/_methylvi_model.py b/src/scvi/external/methylvi/_methylvi_model.py
@@ -157,16 +157,14 @@ def setup_mudata(
 
         Examples
         --------
-        MethylVI.setup_mudata(
-            mdata,
-            mc_layer="mc",
-            cov_layer="cov",
-            batch_key="Platform",
-            methylation_modalities=['mCG', 'mCH'],
-            modalities={
-                "batch_key": "mCG"
-            },
-        )
+        >>> MethylVI.setup_mudata(
+        ...     mdata,
+        ...     mc_layer="mc",
+        ...     cov_layer="cov",
+        ...     batch_key="Platform",
+        ...     methylation_modalities=["mCG", "mCH"],
+        ...     modalities={"batch_key": "mCG"},
+        ... )
 
         """
         if modalities is None:
diff --git a/src/scvi/external/mrvi/__init__.py b/src/scvi/external/mrvi/__init__.py
@@ -1,3 +1,5 @@
+from scvi.external.mrvi_torch._module import TorchMRVAE as MRVAE
+
 from ._model import MRVI
 
-__all__ = ["MRVI"]
+__all__ = ["MRVI", "MRVAE"]
diff --git a/src/scvi/external/poissonvi/_model.py b/src/scvi/external/poissonvi/_model.py
@@ -306,6 +306,7 @@ def differential_accessibility(
         Returns
         -------
         Differential accessibility DataFrame with the following columns:
+
         prob_da
             the probability of the region being differentially accessible
         is_da_fdr
diff --git a/src/scvi/external/resolvi/_module.py b/src/scvi/external/resolvi/_module.py
@@ -100,7 +100,7 @@ class RESOLVAEModel(PyroModule):
     background_ratio:
         Prior for background_proportion
         Equals Gamma(prior_proportions_rate,
-                     prior_proportions_rate/(10*background_ratio*prior_true_amount))
+        prior_proportions_rate/(10*background_ratio*prior_true_amount))
         Default is 0.1
     prior_proportions_rate:
         Rate parameter for the prior proportions.
@@ -1140,7 +1140,7 @@ class RESOLVAE(PyroBaseModuleClass):
     background_ratio:
         Prior for background_proportion
         Equals Gamma(prior_proportions_rate,
-                     prior_proportions_rate/(10*background_ratio*prior_true_amount))
+        prior_proportions_rate/(10*background_ratio*prior_true_amount))
         Default is 0.1
     prior_proportions_rate:
         Rate parameter for the prior proportions.
diff --git a/src/scvi/external/resolvi/_utils.py b/src/scvi/external/resolvi/_utils.py
@@ -147,7 +147,7 @@ def get_normalized_expression_importance(
         return_numpy
             Return a :class:`~numpy.ndarray` instead of a :class:`~pandas.DataFrame`. DataFrame
             includes gene names as columns. If either `n_samples=1` or `return_mean=True`, defaults
-             to `False`. Otherwise, it defaults to `True`.
+            to `False`. Otherwise, it defaults to `True`.
         library_scaling
             If `True`, multiplies the decoded expression by the library size.
         size_scaling
@@ -317,7 +317,7 @@ def get_normalized_expression(
         return_numpy
             Return a :class:`~numpy.ndarray` instead of a :class:`~pandas.DataFrame`. DataFrame
             includes gene names as columns. If either `n_samples=1` or `return_mean=True`, defaults
-             to `False`. Otherwise, it defaults to `True`.
+            to `False`. Otherwise, it defaults to `True`.
         %(de_silent)s
         **kwargs
             Additional keyword arguments passed
@@ -476,7 +476,7 @@ def get_neighbor_abundance(
         return_numpy
             Return a :class:`~numpy.ndarray` instead of a :class:`~pandas.DataFrame`. DataFrame
             includes gene names as columns. If either `n_samples=1` or `return_mean=True`, defaults
-             to `False`. Otherwise, it defaults to `True`.
+            to `False`. Otherwise, it defaults to `True`.
         kwargs
             Additional keyword arguments that have no effect and only serve for compatibility.
 
diff --git a/src/scvi/external/totalanvi/__init__.py b/src/scvi/external/totalanvi/__init__.py
@@ -1,3 +1,4 @@
 from ._model import TOTALANVI
+from ._module import TOTALANVAE
 
-__all__ = ["TOTALANVI"]
+__all__ = ["TOTALANVI", "TOTALANVAE"]
diff --git a/src/scvi/model/__init__.py b/src/scvi/model/__init__.py
@@ -14,6 +14,7 @@
 from ._scanvi import SCANVI
 from ._scvi import SCVI
 from ._totalvi import TOTALVI
+from ._utils import get_max_epochs_heuristic
 
 __all__ = [
     "SCVI",
@@ -27,6 +28,7 @@
     "MULTIVI",
     "AmortizedLDA",
     "utils",
+    "get_max_epochs_heuristic",
 ]
 
 
diff --git a/src/scvi/model/_multivi.py b/src/scvi/model/_multivi.py
@@ -127,7 +127,7 @@ class MULTIVI(
     >>> vae.train()
 
     Notes (for using setup_anndata)
-    -----
+    ---------------------------------
     As of SCVI-Tools v1.4 there is no longer support for setup_anndata for multivi.
     Please use setup_mudata instead.
     """
@@ -771,6 +771,7 @@ def differential_accessibility(
         Returns
         -------
         Differential accessibility DataFrame with the following columns:
+
         prob_da
             the probability of the region being differentially accessible
         is_da_fdr
diff --git a/src/scvi/model/_peakvi.py b/src/scvi/model/_peakvi.py
@@ -451,6 +451,7 @@ def differential_accessibility(
         Returns
         -------
         Differential accessibility DataFrame with the following columns:
+
         prob_da
             the probability of the region being differentially accessible
         is_da_fdr
diff --git a/src/scvi/model/base/_differential.py b/src/scvi/model/base/_differential.py
@@ -191,7 +191,7 @@ def get_bayes_factors(
         delta
             specific case of region inducing differential expression.
             In this case, we suppose that :math:`R \setminus [-\delta, \delta]` does not induce
-             a differential expression (LFC case). If the provided value is `None`, then a proper
+            a differential expression (LFC case). If the provided value is `None`, then a proper
             threshold is determined from the distribution of LFCs across genes.
         pseudocounts
             pseudocount offset used for the mode `change`.
diff --git a/src/scvi/module/_mlxvae.py b/src/scvi/module/_mlxvae.py
@@ -1,12 +1,17 @@
+from __future__ import annotations
+
 import random
-from typing import Any
+from typing import TYPE_CHECKING
 
 import mlx.core as mx
 import mlx.nn as nn
 
 from scvi import REGISTRY_KEYS
 from scvi.module.base import LossOutput
 
+if TYPE_CHECKING:
+    from typing import Any
+
 
 def _lgamma(x: mx.array) -> mx.array:
     """Log-gamma function via Numerical Recipes Lanczos approximation.
diff --git a/src/scvi/module/base/_base_module.py b/src/scvi/module/base/_base_module.py
@@ -368,19 +368,17 @@ def list_obs_plate_vars(self):
         """Model annotation for minibatch training with pyro plate.
 
         A dictionary with:
+
         1. "name" - the name of observation/minibatch plate;
         2. "in" - indexes of model args to provide to encoder network when using amortised
-            inference;
-        3. "sites" - dictionary with
-            keys - names of variables that belong to the observation plate (used to recognise
-            and merge posterior samples for minibatch variables)
-            values - the dimensions in the non-plate axis of each variable (used to construct
-            output layer of encoder network when using amortised inference)
+           inference;
+        3. "sites" - dictionary with keys as names of variables that belong to the observation
+           plate and values as the dimensions in the non-plate axis of each variable.
         """
         return {"name": "", "in": [], "sites": {}}
 
     def on_load(self, model, **kwargs):
-        """Callback function run in :method:`~scvi.model.base.BaseModelClass.load`.
+        """Callback function run in :meth:`~scvi.model.base.BaseModelClass.load`.
 
         For some Pyro modules with AutoGuides, run one training step prior to loading state dict.
         """
diff --git a/src/scvi/nn/__init__.py b/src/scvi/nn/__init__.py
@@ -1,3 +1,5 @@
+from torch.nn.functional import one_hot
+
 from ._base_components import (
     Decoder,
     DecoderSCVI,
@@ -22,4 +24,5 @@
     "MultiEncoder",
     "MultiDecoder",
     "Embedding",
+    "one_hot",
 ]
diff --git a/src/scvi/train/_trainingplans.py b/src/scvi/train/_trainingplans.py
@@ -1967,19 +1967,8 @@ def forward(self, *args, **kwargs):
             pass
 
 
-try:
-    import mlx.core as mx
-    import mlx.nn as nn
-    import mlx.optimizers as optim
-
-    _mlx_available = True
-except (ImportError, OSError):
-    _mlx_available = False
-
-if _mlx_available:
-    import logging
-
-    logger = logging.getLogger(__name__)
+if is_package_installed("mlx"):
+    import mlx
 
     class MlxTrainingPlan(TrainingPlan):
         """Training plan for MLX modules.
@@ -2021,7 +2010,7 @@ def __init__(
             **loss_kwargs,
         ):
             # Create MLX optimizer (AdamW includes weight decay)
-            self.optimizer = optim.AdamW(
+            self.optimizer = mlx.optimizers.AdamW(
                 learning_rate=lr,
                 weight_decay=weight_decay,
                 eps=eps,
@@ -2087,9 +2076,7 @@ def _set_module_training(self, is_training):
             elif hasattr(self.module, "_training"):
                 self.module._training = is_training
             else:
-                logger.warning(
-                    "Unable to set module training state, may affect training performance"
-                )
+                Warning("Unable to set module training state, may affect training performance")
 
         def get_kl_weight(self) -> float:
             """Compute the KL weight for the current step or epoch.
@@ -2135,7 +2122,7 @@ def train_step(self, batch_dict):
             -------
             dict
                 Dictionary containing loss and component metric arrays (not
-                yet evaluated — call ``mx.eval`` on the returned arrays to
+                yet evaluated — call ``mlx.core.eval`` on the returned arrays to
                 materialise them on the GPU).
             """
             # Ensure module is in training mode (validate_step switches to eval)
@@ -2146,26 +2133,27 @@ def train_step(self, batch_dict):
 
             # Convert batch data to MLX arrays
             mlx_batch = {
-                k: mx.array(v) if isinstance(v, np.ndarray) else v for k, v in batch_dict.items()
+                k: mlx.core.array(v) if isinstance(v, np.ndarray) else v
+                for k, v in batch_dict.items()
             }
 
             # Compute KL weight
             kl_weight = self.get_kl_weight()
 
             # Single forward+backward: value_and_grad with aux returns
             # (loss, aux), grads — we get LossOutput in one pass
-            loss_and_grad_fn = nn.value_and_grad(self.module, self._loss_fn)
+            loss_and_grad_fn = mlx.nn.value_and_grad(self.module, self._loss_fn)
             (loss, loss_output), grads = loss_and_grad_fn(self.module, mlx_batch, kl_weight)
 
             # Global-norm gradient clipping (vectorised, no per-tensor loops)
-            grads, _ = optim.clip_grad_norm(grads, max_norm=5.0)
+            grads, _ = mlx.optimizers.clip_grad_norm(grads, max_norm=5.0)
 
             # Update parameters
             self.optimizer.update(self.module, grads)
 
-            # Single mx.eval at the end — materialises the whole lazy graph
+            # Single mlx.core.eval at the end — materialises the whole lazy graph
             # (parameters, optimizer state, and the metrics we want to read)
-            mx.eval(
+            mlx.core.eval(
                 self.module.parameters(),
                 self.optimizer.state,
                 loss,
@@ -2203,15 +2191,16 @@ def validate_step(self, batch_dict):
 
             # Convert batch data to MLX arrays
             mlx_batch = {
-                k: mx.array(v) if isinstance(v, np.ndarray) else v for k, v in batch_dict.items()
+                k: mlx.core.array(v) if isinstance(v, np.ndarray) else v
+                for k, v in batch_dict.items()
             }
 
             kl_weight = self.get_kl_weight()
 
             _, _, loss_output = self.module(mlx_batch, kl_weight=kl_weight)
 
             # Single eval for all metrics
-            mx.eval(
+            mlx.core.eval(
                 loss_output.loss,
                 loss_output.reconstruction_loss_sum,
                 loss_output.kl_local_sum,
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/dataloaders/sparse_utils.py b/tests/dataloaders/sparse_utils.py
diff --git a/tests/dataloaders/test_dataloaders.py b/tests/dataloaders/test_dataloaders.py
diff --git a/tests/model/base/test_base_model.py b/tests/model/base/test_base_model.py
