Foundations · INTEGRATION

Masked sums, means and weighted reductions

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Reduce a tensor over selected axes while excluding masked entries. Return either the weighted sum or the weighted mean; a mean with no valid weight is exactly zero.

Function signature

masked_reductions(
    x: torch.Tensor,
    mask: torch.Tensor,
    weights: torch.Tensor | None = None,
    dim: int | list[int] | tuple[int, ...] | None = -1,
    keepdim: bool = False,
    reduction: str = 'mean',
    eps: float = 1e-08,
) -> torch.Tensor

Worked example

The same slice, two reductions

Take x=[2,6,100], mask=[True,True,False], weights=[1,3,9], dim=0 and eps=0. The included weighted sum is 1×2+3×6=20; the included total weight is 1+3=4. reduction="sum" returns scalar 20; reduction="mean" returns scalar 5. The masked 100 contributes neither value nor weight.

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[d_0,\ldots,d_{R-1}]$ with arbitrary rank $R\ge1$ .Values to reduce at included positions.
`mask`torch.Tensor	Any shape broadcastable to `x.shape`.True includes an entry; False excludes it.
`weights`torch.Tensor \| NoneDefault `None`	Any shape broadcastable to `x.shape`.Optional nonnegative contribution weights; None gives each included entry weight one.
`dim`int \| list[int] \| tuple[int, ...] \| NoneDefault `-1`	One axis, a nonempty list/tuple of distinct axes, or None for all axes. Negative indices count from the end.
`keepdim`boolDefault `False`	Keep every reduced axis with size one when True; otherwise remove those axes and preserve the order of the others.
`reduction`strDefault `'mean'`	Exactly "mean" or "sum". This exercise does not support an unreduced "none" result.
`eps`floatDefault `1e-08`	Nonnegative offset added once to the summed valid weight in a mean denominator. It is ignored for sums; a zero-weight mean is exactly zero.

Return

torch.Tensor · The shape of x with the axes in dim removed, or replaced by 1 when keepdim=True. dim=None reduces all axes.

Return Weighted sum or mean; reduced axes removed or size one according to keepdim.

Required behavior

Include entries where mask is True. Broadcast both mask and optional weights to x.shape before reducing; absent weights mean one at every entry.
For each retained slice compute $S=\sum_i w_i x_i$ and $D=\sum_i w_i$ over included entries. A broadcast weight contributes once for every expanded entry.
Return $S$ for sum. For mean, return $S/(D+\varepsilon)$ when $D>0$ , and exactly zero when $D=0$ . Add epsilon once after summing weights.
Use only the selected axes. Remove them unless keepdim=True, in which case each has size one. A full reduction without keepdim returns a scalar tensor, shape $()$ .
dim is one integer axis, a nonempty list/tuple of distinct axes, or None for all axes. Negative axes count from the end. reduction supports only "sum" and "mean".
All floating inputs are finite and share dtype/device. Weights are nonnegative; epsilon is finite and nonnegative. Masked padding may be arbitrary finite data.
Empty reduced slices return zero; empty retained axes remain empty. A zero-total-weight mean has zero gradients. A weighted sum retains its ordinary weight gradients even at zero weights.
The mean is global over the requested domain, not an average of smaller unweighted means. Preserve strided inputs and do not modify the supplied tensors.

Why this matters

A weighted mean divides the total weighted value by the total expanded weight; averaging smaller means can give the wrong weighting.

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_reductions(
    x: torch.Tensor,  # (d_0,...,d_{R-1}) with arbitrary rank R>=1.
    mask: torch.Tensor,  # Any shape broadcastable to x.shape.
    weights: torch.Tensor | None = None,  # Any shape broadcastable to x.shape.
    dim: int | list[int] | tuple[int, ...] | None = -1,
    keepdim: bool = False,
    reduction: str = 'mean',
    eps: float = 1e-08,
) -> torch.Tensor:
    """Return a tensor. The shape of x with the axes in dim removed, or replaced by 1 when
        keepdim=True. dim=None reduces all axes."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement masked_reductions")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Related exercises

Sources and further reading

aqlaboratory/openfold · tensor_utils.py

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.