Stable masked Poisson objective for genomic counts

Your task

Function signature

masked_poisson_deviance(
    counts: torch.Tensor,
    rate: torch.Tensor,
    mask: torch.Tensor,
    objective: str = 'deviance',
    reduction: str = 'mean',
    eps: float = 1e-06,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Zero-count and perfectly predicted entries

import torch

counts = torch.tensor([[[0.0, 2.0], [3.0, 0.0]]], dtype=torch.float64)
rate = torch.tensor([[[0.875, 1.875], [2.875, 3.875]]], dtype=torch.float64, requires_grad=True)
mask = torch.tensor([[[True, True], [True, True]]], dtype=torch.bool)
objective = 'deviance'
reduction = 'mean'
eps = 0.125

result = masked_poisson_deviance(counts=counts, rate=rate, mask=mask, objective=objective, reduction=reduction, eps=eps)

Expected return

(torch.tensor(2.5, dtype=torch.float64), torch.tensor([1.0, 4.0], dtype=torch.float64))

Parameters

Return

tuple[torch.Tensor, torch.Tensor]

Required behavior

• Compute a masked Poisson objective and a per-track reduction with stable zero-count behavior. counts and rate have identical shape [B,L,T], with B,T>0 and L\ge 0. Floating inputs share dtype/device (float32 or float64); reductions preserve that dtype. mask is bool and broadcastable to [B,L,T], with True means observed.
• For example use [B,L,1] for a positional mask, [1,1,T] for a track mask, or a scalar bool. At observed positions \mathrm{counts}\ge 0 and \mathrm{rate}\ge 0 must be finite. Missing payloads may be arbitrary, including negative values, NaN, or infinity: replace them before any logarithm/arithmetic.
• Do not infer missingness from a count of zero. options include eps, written as \varepsilon (positive float, default 10^{-6}), objective ('deviance', default, or 'nll'), and reduction ('mean', default, or 'sum'). The effective mean is \mu=\mathrm{rate}+\varepsilon; \varepsilon is additive, not a clamp and not added to counts.
• For each observed entry compute either deviance: d=2\cdot (\mu-\mathrm{counts}+\mathrm{counts}\cdot (\operatorname{log}(\mathrm{counts})-\operatorname{log}(\mu))), nll: d=\mu-\mathrm{counts}\cdot \operatorname{log}(\mu). Define \mathrm{counts}\cdot \operatorname{log}(\mathrm{counts})=0 when \mathrm{counts}=0.
• Make the log-count argument safe before selecting a branch, so even unused branches have finite derivatives. The nll option omits the target-only \operatorname{log}(\mathrm{counts}!) constant; deviance includes the conventional factor 2 and target-optimal subtraction, not a source-specific half/shifted deviance.
• Missing entries contribute zero, not \varepsilon. Return loss, a scalar reduction over B,L,T, and per_track [T], reducing over B,L for each track.
• For mean, divide each sum by its own number of valid elements after broadcasting the mask. Global loss is not the unweighted mean of per-track means. Any empty/all-masked reduction is exactly zero (denominator at least 1), including L=0. No residual, regularization, or learned weights are added.
• A zero-count observed entry still contributes \mu (nll) or 2\cdot \mu (deviance). Negative observed targets are outside the input contract.

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 masked_poisson_deviance(
    counts: torch.Tensor,  # (B, L, T)
    rate: torch.Tensor,  # (B, L, T)
    mask: torch.Tensor,  # Broadcastable to (B, L, T)
    objective: str = 'deviance',
    reduction: str = 'mean',
    eps: float = 1e-06,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    loss: Scalar tensor, shape ().
    per_track: (T,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement masked_poisson_deviance")

Sources and further reading

• google-deepmind/alphagenome_research · losses.py
• calico/baskerville · metrics.py