implement utils in torch

Author: lordy5

src/scvi/external/mrvi/_model.py

@@ -4,7 +4,6 @@
 import warnings
 from typing import TYPE_CHECKING
 
-import jax
 import numpy as np
 import torch
 import xarray as xr
@@ -1335,12 +1334,11 @@ def h_inference_fn(extra_eps, batch_index_cf, batch_offset_eps):
                     )
                 # batch_offset shape (mc_samples, n_batch, n_cells, n_latent)
 
-                f_ = jax.vmap(
+                f_ = torch.vmap(
                     h_inference_fn, in_axes=(0, None, 0), out_axes=0
                 )  # fn over MC samples
-                f_ = jax.vmap(f_, in_axes=(1, None, None), out_axes=1)  # fn over covariates
-                f_ = jax.vmap(f_, in_axes=(None, 0, 1), out_axes=0)  # fn over batches
-                h_fn = jax.jit(f_)
+                f_ = torch.vmap(f_, in_axes=(1, None, None), out_axes=1)  # fn over covariates
+                h_fn = torch.vmap(f_, in_axes=(None, 0, 1), out_axes=0)  # fn over batches
 
                 x_1 = h_fn(betas_covariates, batch_index_, betas_offset_)
                 x_0 = h_fn(betas_null, batch_index_, betas_offset_)
@@ -1491,7 +1489,7 @@ def _construct_design_matrix(
         add_batch_specific_offsets: bool,
         store_lfc: bool,
         store_lfc_metadata_subset: list[str] | None = None,
-    ) -> tuple[jax.Array, npt.NDArray, jax.Array, jax.Array | None]:
+    ) -> tuple[torch.Tensor, npt.NDArray, torch.Tensor, torch.Tensor | None]:
         """Construct a design matrix of samples and covariates.
 
         Starting from a list of sample covariate keys, construct a design matrix of samples and

src/scvi/external/mrvi/_module.py

@@ -3,8 +3,6 @@
 import warnings
 from typing import TYPE_CHECKING
 
-# import flax.linen as nn
-import jax
 import torch
 import torch.distributions as dist
 import torch.nn as nn
@@ -34,9 +32,6 @@
 }
 DEFAULT_QU_KWARGS = {}
 
-# Lower stddev leads to better initial loss values
-_normal_initializer = jax.nn.initializers.normal(stddev=0.1)
-
 
 class DecoderZXAttention(nn.Module):
     """Attention-based decoder.

src/scvi/external/mrvi/_utils.py

@@ -2,13 +2,10 @@
 
 from typing import TYPE_CHECKING
 
-from jax import jit
-
 from scvi.external.mrvi._types import _ComputeLocalStatisticsRequirements
 
 if TYPE_CHECKING:
-    from jax import Array
-    from jax.typing import ArrayLike
+    from torch import Tensor
 
     from scvi.external.mrvi._types import MRVIReduction
 
@@ -58,17 +55,16 @@ def _parse_local_statistics_requirements(
     )
 
 
-@jit
-def rowwise_max_excluding_diagonal(matrix: ArrayLike) -> Array:
+def rowwise_max_excluding_diagonal(matrix: Tensor) -> Tensor:
     """Get the rowwise maximum of a matrix excluding the diagonal."""
-    import jax.numpy as jnp
+    import torch
 
     assert matrix.ndim == 2
     num_cols = matrix.shape[1]
-    mask = (1 - jnp.eye(num_cols)).astype(bool)
-    return (jnp.where(mask, matrix, -jnp.inf)).max(axis=1)
+    mask = (1 - torch.eye(num_cols)).astype(bool)
+    return (torch.where(mask, matrix, -torch.inf)).max(axis=1)
 
 
-def simple_reciprocal(w: ArrayLike, eps: float = 1e-6) -> Array:
+def simple_reciprocal(w: Tensor, eps: float = 1e-6) -> Tensor:
     """Convert distances to similarities via a reciprocal."""
     return 1.0 / (w + eps)