TOTALVAE.inference(x, y, batch_index=None, label=None, n_samples=1, cont_covs=None, cat_covs=None)[source]

Internal helper function to compute necessary inference quantities.

We use the dictionary px_ to contain the parameters of the ZINB/NB for genes. The rate refers to the mean of the NB, dropout refers to Bernoulli mixing parameters. scale refers to the quanity upon which differential expression is performed. For genes, this can be viewed as the mean of the underlying gamma distribution.

We use the dictionary py_ to contain the parameters of the Mixture NB distribution for proteins. rate_fore refers to foreground mean, while rate_back refers to background mean. scale refers to foreground mean adjusted for background probability and scaled to reside in simplex. back_alpha and back_beta are the posterior parameters for rate_back. fore_scale is the scaling factor that enforces rate_fore > rate_back.

px_["r"] and py_["r"] are the inverse dispersion parameters for genes and protein, respectively.

x : TensorTensor

tensor of values with shape (batch_size, n_input_genes)

y : TensorTensor

tensor of values with shape (batch_size, n_input_proteins)

batch_index : Tensor | NoneOptional[Tensor] (default: None)

array that indicates which batch the cells belong to with shape batch_size

label : Tensor | NoneOptional[Tensor] (default: None)

tensor of cell-types labels with shape (batch_size, n_labels)


Number of samples to sample from approximate posterior


Continuous covariates to condition on


Categorical covariates to condition on

Return type

{str: Tensor | {str: Tensor}}Dict[str, Union[Tensor, Dict[str, Tensor]]]