netmap.model.zinbautoencoder.ZINBAutoencoder

class netmap.model.zinbautoencoder.ZINBAutoencoder(*args: Any, **kwargs: Any)

Bases: Module

Autoencoder with three decoder heads for ZINB-distributed scRNA-seq data.

The shared encoder compresses gene expression into a latent space. Three independent decoder branches predict the ZINB parameters:

• mu (mean expression) via Softplus activation.

• theta (dispersion) via Softplus activation.

• pi (zero-inflation probability) via Softplus activation.

Four mutually exclusive forward-mode flags redirect forward() to return a single output tensor for Captum attribution:

• forward_mu_only: return only mu.

• forward_theta_only: return only theta.

• latent_only: return the latent embedding.

• forward_pi_only: return only pi.

When all flags are False (default), forward returns (mu, theta, pi).

Parameters:

• input_dim (int) – Number of input genes.

• latent_dim (int) – Dimensionality of the latent space.

• dropout_rate (float) – Dropout probability after each hidden layer. Defaults to 0.0.

• hidden_dims (list of int) – Width of each hidden layer in the encoder; all three decoders mirror this in reverse. Defaults to [64].

__init__(input_dim, latent_dim, dropout_rate=0.0, hidden_dims=[64])

compute_loss(x)

Compute the ZINB loss for a batch.

Parameters:

x (torch.Tensor) – Input count matrix of shape (batch, input_dim).

Returns:

Scalar ZINB negative log-likelihood loss.

Return type:

torch.Tensor

forward(x)

Run a forward pass, returning output controlled by the mode flags.

When all mode flags are False, returns (mu, theta, pi). Set one of forward_mu_only, forward_theta_only, latent_only, or forward_pi_only to True before passing the model to a Captum explainer so that it returns a single tensor.

Parameters:

x (torch.Tensor) – Input gene expression tensor of shape (n_cells, input_dim).

Returns:

Depending on the active flag —

mu, theta, latent, pi, or (mu, theta, pi) tuple.

Return type:

torch.Tensor or tuple