Library-size-scaled rates and batch-specific dispersion

Your task

Function signature

library_size_dispersion_decoder(
    gene_logits: torch.Tensor,
    gene_mask: torch.Tensor,
    library: torch.Tensor,
    log_dispersion: torch.Tensor,
    batch_id: torch.Tensor | None = None,
    library_mode: str = 'log',
    dispersion_mode: str = 'gene',
) -> dict[str, torch.Tensor]

Worked example

Allocate a library to unequal gene proportions

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Build a masked gene composition, library-scaled RNA rates and positive inverse dispersion. gene_logits and gene_mask have shape [B,G], with cells on axis 0 and genes on axis 1. gene_mask is boolean and True means a gene is in this cell's composition support; it is not a mask inferred from observed nonzero counts.
• library has shape [B,1]. log_dispersion is [G] for gene mode or [K,G] for batch mode. batch_id is integer [B]. B,G,K \ge 1. Inputs are finite, and all tested library/dispersion exponentials fit the dtype. No clipping, softplus-scale alternative or \epsilon is used.
• Choose the library and dispersion conventions: library_mode="log" (default): \mathrm{depth}=\operatorname{exp}(\mathrm{library}). The supplied log depth may itself originate from a latent sample or an observed library; this operation does not infer it.
• library_mode="observed": library is a nonnegative raw molecule count/size factor, so \mathrm{depth}=\mathrm{library}; zero is allowed and produces zero rates. dispersion_mode="gene" (default): \theta[b,g]=\operatorname{exp}(\mathrm{log\_dispersion}[g]); batch_id is ignored.
• dispersion_mode="batch": \theta[b,g]=\operatorname{exp}(\mathrm{log\_dispersion}[\mathrm{batch\_id}[b],g]); each batch_id must be in [0,K).
• For every nonempty mask row, proportions [b,g] is the softmax of gene_logits over only True genes, and zero at False genes. Implement the softmax manually with a masked row maximum and safe denominator. library_size is the effective depth: depth for nonempty support and zero for an all-masked row.
• \mathrm{rate} = \mathrm{library\_size}\cdot \mathrm{proportions}. Thus \sum_{g} \mathrm{rate}[b,g] = \mathrm{library\_size}[b,0]; it equals the supplied raw/exp depth whenever there is at least one supported gene. All-masked rows have exactly zero proportions, effective depth and rates even when supplied depth was positive.
• Dispersion stays positive at every gene, including masked genes, since it is a parameter table rather than an observation. Return new parameter tensors, not residual updates, and do not modify inputs. Reject unknown modes, negative observed depth, and out-of-range batch IDs with ValueError.
• Other declared shapes/dtypes/finiteness are preconditions.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def library_size_dispersion_decoder(
    gene_logits: torch.Tensor,  # (B, G)
    gene_mask: torch.Tensor,  # (B, G)
    library: torch.Tensor,  # (B, 1)
    log_dispersion: torch.Tensor,  # (G,) or (K, G)
    batch_id: torch.Tensor | None = None,  # (B,)
    library_mode: str = 'log',
    dispersion_mode: str = 'gene',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    proportions: (B, G).
    rate: (B, G).
    inverse_dispersion: (B, G).
    library_size: (B, 1)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement library_size_dispersion_decoder")

Sources and further reading

• scverse/scvi-tools · _vae.py
• scverse/scvi-tools · _base_components.py
• scverse/scvi-tools · _base_components.py