torchcell.losses.distributional.masked_mean

masked_mean(elem: Tensor, feature_mask: Tensor | None) → Tensor

Mean of elem over the entries feature_mask selects.

WHY THIS EXISTS. The masked-label objective scores a prediction ONLY on the genes that are still hidden at unmasking step \(k\). Genes whose true value has been revealed are model INPUT at that step, so including them in the loss lets the network copy the input to the output – train loss collapses, train Pearson inflates, and none of it transfers to validation, where nothing is revealed. Restricting the reduction is the whole correctness requirement of that objective.

feature_mask is [B, F] (True = score this entry) while elem may be [B, F] (point) or [B, F, K] (quantile / distributional), so the mask is broadcast along any trailing axes – a quantile knot is scored exactly when its gene is.

Parameters:

• elem – Elementwise loss [B, F] or [B, F, ...].

• feature_mask – Optional bool [B, F]; None reduces over everything.

Returns:

Scalar mean over selected entries. Returns a graph-connected zero when the mask selects nothing, so DDP find-unused stays happy.