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, parent=<flax.linen.module._Sentinel object>, name=None)

Bases: JaxBaseModuleClass

Variational autoencoder model.

Attributes table

dropout_rate
eps
gene_likelihood
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)	Run generative model.
inference(x[, n_samples])	Run inference model.
loss(tensors, inference_outputs, ...[, ...])	Compute loss.
setup()	Setup model.

Attributes

JaxVAE.dropout_rate: Tunable_[float] = 0.0

JaxVAE.eps: Tunable_[float] = 1e-08

JaxVAE.gene_likelihood: Tunable_[str] = 'nb'

JaxVAE.n_hidden: Tunable_[int] = 128

JaxVAE.n_latent: Tunable_[int] = 30

JaxVAE.n_layers: Tunable_[int] = 1

JaxVAE.name: Optional[str] = None

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

JaxVAE.required_rngs

JaxVAE.scope: Optional[Scope] = None

JaxVAE.training: bool = True

JaxVAE.n_input: int

JaxVAE.n_batch: int

Methods

JaxVAE.generative(x, z, batch_index)

Run generative model.

Return type:

dict

JaxVAE.inference(x, 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.