Keep missing, masked, padded, and truly zero values separate

Your task

Function signature

expression_masks_and_masked_objectives(
    values: torch.Tensor,
    predictions: torch.Tensor,
    targets: torch.Tensor,
    nonzero_logits: torch.Tensor,
    masks: dict,
    uniform: torch.Tensor | None = None,
    selection_mode: str = 'provided',
    mask_probability: float = 0.15,
    mask_value: float = -1.0,
    nonzero_weight: float = 1.0,
) -> dict[str, torch.Tensor]

Worked example

Measured zero is still a target

import torch

values = torch.tensor([[0.0, 2.0, 0.0, -9.0]], dtype=torch.float64, requires_grad=True)
predictions = torch.tensor([[1.0, 0.0, 99.0, 99.0]], dtype=torch.float64, requires_grad=True)
targets = torch.tensor([[0.0, 2.0, 0.0, -9.0]], dtype=torch.float64, requires_grad=True)
nonzero_logits = torch.tensor([[0.0, 0.0, 0.0, 0.0]], dtype=torch.float64, requires_grad=True)
uniform = torch.tensor([[0.1, 0.2, 0.3, 0.4]], dtype=torch.float64)
masks = {
    'valid_mask': torch.tensor([[True, True, False, True]], dtype=torch.bool),
    'observed_mask': torch.tensor([[True, True, False, False]], dtype=torch.bool),
    'eligible_mask': torch.tensor([[True, True, True, True]], dtype=torch.bool),
    'requested_mask': torch.tensor([[True, True, True, True]], dtype=torch.bool),
}
selection_mode = 'provided'
mask_probability = 0.15
mask_value = -1.0
nonzero_weight = 1.0

result = expression_masks_and_masked_objectives(values=values, predictions=predictions, targets=targets, nonzero_logits=nonzero_logits, uniform=uniform, masks=masks, selection_mode=selection_mode, mask_probability=mask_probability, mask_value=mask_value, nonzero_weight=nonzero_weight)

Expected return

{
    'corrupted_values': torch.tensor([[-1.0, -1.0, 0.0, -9.0]], dtype=torch.float64),
    'selected_mask': torch.tensor([[True, True, False, False]], dtype=torch.bool),
    'observed_zero_mask': torch.tensor([[True, False, False, False]], dtype=torch.bool),
    'missing_mask': torch.tensor([[False, False, False, True]], dtype=torch.bool),
    'padding_mask': torch.tensor([[False, False, True, False]], dtype=torch.bool),
    'selected_count': torch.tensor(2, dtype=torch.int64),
    'mse': torch.tensor(2.5, dtype=torch.float64),
    'nonzero_nll': torch.tensor(0.6931471805599453, dtype=torch.float64),
    'loss': torch.tensor(3.1931471805599454, dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Corrupt eligible expression positions using explicit masks and supplied selection decisions, then compute masked regression and nonzero classification losses. All value and mask tensors have shape [B,L], cells then gene-token slots; B,L \ge 1.
• values are the model's input values, targets are the original expression targets, predictions are real-valued reconstructions, and nonzero_logits are Bernoulli logits for P(\mathrm{target}>0), not a structural-zero probability. Values are finite; observed expression targets are nonnegative.
• A measured zero target is valid supervision. Missing/padded entries can carry any finite sentinel. Mask polarity is always True to include: valid_mask means a nonpadding slot, observed_mask means an available target measurement, and eligible_mask allows corruption (set False to protect special tokens or other protected positions).
• \mathrm{pool}=\mathrm{valid\_mask} \land \mathrm{observed\_mask} \land \mathrm{eligible\_mask}. Never infer these masks from values equaling zero or a sentinel. options: * selection_mode="provided" (default): \mathrm{chosen}=\mathrm{requested\_mask}, a supplied bool [B,L].
• * selection_mode="random": \mathrm{chosen}=(\mathrm{uniform} < \mathrm{mask\_probability}), where uniform is supplied [B,L] in [0,1); do not sample randomness. mask_probability defaults to 0.15 and is in [0,1]. This is independent Bernoulli corruption, not an exact fixed number per cell.
• requested_mask is ignored in this mode; uniform is ignored in provided mode. * mask_value (default -1.0) is the finite replacement value at \mathrm{selected}=\mathrm{pool} \land \mathrm{chosen}; every other input slot, including padding, is left unchanged in corrupted_values.
• * nonzero_weight (default 1.0) is nonnegative and weights only the Bernoulli term in the combined scalar loss. Let C be the total selected count over the minibatch. \mathrm{mse}=\sum_{\mathrm{selected}}(\mathrm{predictions}-\mathrm{targets})^{2} / C.
• Let y=1[\mathrm{target}>0], t=\mathrm{nonzero\_logits}; \mathrm{nonzero\_nll}=\sum_{\mathrm{selected}}(\operatorname{max}(t,0)-t\cdot y+\mathrm{log1p}(\operatorname{exp}(-\left|t\right|))) / C. \mathrm{loss}=\mathrm{mse}+\mathrm{nonzero\_weight}\cdot \mathrm{nonzero\_nll}.
• This is a global selected-token mean, not a mean of cell means. If C=0, all three losses are exactly zero. Sanitize excluded entries before squared-error/probability arithmetic and use a safe denominator. No \epsilon is required. Return the explicit selected, observed-zero, missing and padding masks plus selected_count.
• \mathrm{observed\_zero\_mask}=\mathrm{valid} \land \mathrm{observed} \land (\mathrm{targets}=0), \mathrm{missing\_mask}=\mathrm{valid} \land \neg \mathrm{observed}, \mathrm{padding\_mask}=\neg \mathrm{valid}; these diagnostic masks are independent of whether a position is eligible or selected. Do not mutate any inputs.
• Reject unknown selection modes, probabilities outside [0,1], and negative nonzero_weight with ValueError. Other declared finite values, shapes and uniform ranges are caller 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 expression_masks_and_masked_objectives(
    values: torch.Tensor,  # (B, L)
    predictions: torch.Tensor,  # (B, L)
    targets: torch.Tensor,  # (B, L)
    nonzero_logits: torch.Tensor,  # (B, L)
    masks: dict,
    uniform: torch.Tensor | None = None,  # (B, L)
    selection_mode: str = 'provided',
    mask_probability: float = 0.15,
    mask_value: float = -1.0,
    nonzero_weight: float = 1.0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    corrupted_values: (B, L).
    selected_mask: (B, L).
    observed_zero_mask: (B, L).
    missing_mask: (B, L).
    padding_mask: (B, L).
    selected_count: ().
    mse: ().
    nonzero_nll: ().
    loss: ()."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement expression_masks_and_masked_objectives")

Sources and further reading

• bowang-lab/scGPT · gene_tokenizer.py
• bowang-lab/scGPT · loss.py
• huggingface.co