Stable negative-binomial log probability

Your task

Function signature

negative_binomial_log_prob(
    counts: torch.Tensor,
    mean: torch.Tensor,
    inverse_dispersion: torch.Tensor,
    observed_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Unit mean and dispersion give a geometric distribution

import torch

counts = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)
mean = torch.tensor([[1.0, 1.0, 1.0, 1.0]], dtype=torch.float64, requires_grad=True)
inverse_dispersion = torch.tensor([1.0, 1.0, 1.0, 1.0], dtype=torch.float64, requires_grad=True)
observed_mask = torch.tensor([[True, True, True, True]], dtype=torch.bool)

result = negative_binomial_log_prob(counts=counts, mean=mean, inverse_dispersion=inverse_dispersion, observed_mask=observed_mask)

Expected return

(torch.tensor([[-0.6931471805599453, -1.3862943611198906, -2.0794415416798357, -2.772588722239781]], dtype=torch.float64), torch.tensor([6.931471805599452], dtype=torch.float64))

Parameters

Return

tuple[torch.Tensor, torch.Tensor]

Required behavior

• Implement an exact, masked negative-binomial log mass from primitive tensor operations. Cells are axis 0 and genes axis 1. counts has shape [B,G] and contains nonnegative integer-valued observations; integer or floating count storage is permitted, but counts are discrete and never differentiated. mean is [B,G].
• inverse_dispersion is [G] (shared across cells) or [B,G]. Both are finite and strictly positive everywhere. observed_mask is boolean [B,G], with True meaning measured. A measured count of zero must contribute its NB log probability; False entries are missing and may contain an arbitrary finite count sentinel. B,G \ge 1.
• For each observed (b,g), compute: \log p(x)=\log\Gamma(x+\theta)-\log\Gamma(\theta)-\log\Gamma(x+1)+\theta\log\frac{\theta}{\theta+\mu}+x\log\frac{\mu}{\theta+\mu} . Return zero at missing entries, and \texttt{cell\_nll}[b]=-\sum_g\texttt{log\_prob}[b,g]. An entirely missing cell therefore has \mathrm{cell\_nll}=0.
• Reject negative observed counts and nonpositive mean or inverse dispersion with ValueError. Other declared shapes, finite values and count integrality are caller preconditions. No \epsilon is added: this is the exact positive-parameter NB, not an \epsilon-perturbed source formula.
• Use stable log-ratios rather than exponentiating huge logits or forming \theta+\mu when avoidable. If a=\operatorname{log}(\mu)-\operatorname{log}(\theta), \operatorname{log}(1+\operatorname{exp}(a)) = \operatorname{max}(a,0)+\operatorname{log1p}(\operatorname{exp}(-\left|a\right|)).
• Sanitize missing entries before probability arithmetic, then mask the outputs. There are no options. torch.lgamma is the sole additional permitted numerical operation; torch.distributions is forbidden in candidate code.

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 negative_binomial_log_prob(
    counts: torch.Tensor,  # (B, G)
    mean: torch.Tensor,  # (B, G)
    inverse_dispersion: torch.Tensor,  # (G,) or (B, G)
    observed_mask: torch.Tensor,  # (B, G)
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    log_prob: (B, G).
    cell_nll: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement negative_binomial_log_prob")

Sources and further reading

• scverse/scvi-tools · _negative_binomial.py