scvi.module.JaxVAE

class scvi.module.JaxVAE(n_input, n_batch, n_hidden=128, n_latent=30, dropout_rate=0.0, n_layers=1, gene_likelihood='nb', eps=1e-08, training=True, n_continuous_cov=0, n_cats_per_cov=(), parent=<flax.linen.module._Sentinel object>, name=None)

Bases: JaxBaseModuleClass

Variational autoencoder model.

Attributes table

dropout_rate
eps
gene_likelihood
n_cats_per_cov
n_continuous_cov
n_hidden
n_latent
n_layers
name
parent
required_rngs	Returns a tuple of rng sequence names required for this Flax module.
scope
training
n_input
n_batch

Methods table

generative(x, z, batch_index[, cont_covs, ...])	Run generative model.
inference(x[, cont_covs, cat_covs, n_samples])	Run inference model.
loss(tensors, inference_outputs, ...[, ...])	Compute loss.
setup()	Setup model.

Attributes

JaxVAE.dropout_rate: float = 0.0

JaxVAE.eps: float = 1e-08

JaxVAE.gene_likelihood: str = 'nb'

JaxVAE.n_cats_per_cov: tuple[int, ...] = ()

JaxVAE.n_continuous_cov: int = 0

JaxVAE.n_hidden: int = 128

JaxVAE.n_latent: int = 30

JaxVAE.n_layers: int = 1

JaxVAE.name: Optional[str] = None

JaxVAE.parent: Union[Module, Scope, _Sentinel, None] = None

JaxVAE.required_rngs

JaxVAE.scope: Scope | None = None

JaxVAE.training: bool = True

JaxVAE.n_input: int

JaxVAE.n_batch: int

Methods

JaxVAE.generative(x, z, batch_index, cont_covs=None, cat_covs=None)

Run generative model.

Return type:

dict

JaxVAE.inference(x, cont_covs=None, cat_covs=None, n_samples=1)

Run inference model.

Return type:

dict

JaxVAE.loss(tensors, inference_outputs, generative_outputs, kl_weight=1.0)

Compute loss.

JaxVAE.setup()

Setup model.