Function masked_mean

src/post_training/utils.py:116–121

Mean of ``values`` over positions where ``mask`` is truthy (safe if mask empty).

(values: torch.Tensor, mask: torch.Tensor)

# --- Masked reductions -------------------------------------------------------

def masked_mean(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Mean of ``values`` over positions where ``mask`` is truthy (safe if mask empty)."""
    mask = mask.to(values.dtype)
    total = (values * mask).sum()
    count = mask.sum().clamp(min=1.0)
    return total / count


def masked_mean_per_row(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: