Foundations

Weighted Masked Mean

EasyRaw PyTorchAbout 10 min2 public examples plus private tests

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

Your task

For each index left after reducing the selected axes, return the included weighted sum divided by the included total weight plus eps. Broadcast nonnegative weights to x first; a slice with zero total weight returns zero.

Function signature

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

Worked example

One included row and one empty row

The first row returns (1×2 + 3×4)/(1+3) = 3.5. The third value is masked and contributes neither value nor weight. The entirely masked second row returns zero.

import torch

x = torch.tensor([[2.0, 4.0, 100.0], [6.0, 8.0, 10.0]], dtype=torch.float64, requires_grad=True)
mask = torch.tensor([[True, True, False], [False, False, False]], dtype=torch.bool)
weights = torch.tensor([[1.0, 3.0, 2.0]], dtype=torch.float64, requires_grad=True)
dim = -1
keepdim = False
eps = 0.0

result = weighted_masked_mean(x=x, mask=mask, weights=weights, dim=dim, keepdim=keepdim, eps=eps)

Expected return

torch.tensor([3.5, 0.0], dtype=torch.float64)

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[d_0,\ldots,d_{R-1}]$ Values to average over the selected axes.
`mask`torch.Tensor	Any shape broadcastable to `x.shape`.True includes an entry; broadcasting follows PyTorch rules.
`weights`torch.Tensor	Any shape broadcastable to `x.shape`.Nonnegative weights broadcast to x before both numerator and denominator.
`dim`int \| list[int] \| tuple[int, ...] \| NoneDefault `-1`	Axis to reduce or normalize. Negative indices count from the last axis. None reduces every axis; a list/tuple reduces all listed axes.
`keepdim`boolDefault `False`	Keep each reduced axis with length one when True; otherwise remove it.
`eps`floatDefault `1e-08`	Finite nonnegative value added once to the total included weight. A zero-total-weight slice returns 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

Broadcast mask to exactly x.shape. True includes an entry; False excludes it, regardless of the finite padding value.
dim is an integer, a nonempty list or tuple of distinct axes, or None to reduce all axes. Negative axes count from the end. Preserve the order of the unreduced axes.
With keepdim=True, keep each reduced axis at size one. Otherwise remove it; reducing every axis returns a scalar tensor with shape $()$ .
Empty reduced slices and entirely masked slices return zero. Empty unreduced axes remain empty. Inputs may be strided.
Broadcast both weights and mask to x.shape. Weights are required, finite, and nonnegative.
Over included entries compute $S=\sum_i w_i x_i$ and $D=\sum_i w_i$ . Return $S/(D+\varepsilon)$ for $D>0$ ; otherwise return zero.
eps is finite and nonnegative, default $10^{-8}$ , and is added once after summing weights. Both the numerator and denominator remain differentiable with respect to weights.
A zero-total-weight slice returns zero with zero gradients. Repeated broadcast weights accumulate gradients from every expanded use.

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 weighted_masked_mean(
    x: torch.Tensor,  # (d_0,...,d_{R-1})
    mask: torch.Tensor,  # Any shape broadcastable to x.shape.
    weights: torch.Tensor,  # Any shape broadcastable to x.shape.
    dim: int | list[int] | tuple[int, ...] | None = -1,
    keepdim: bool = False,
    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 weighted_masked_mean")

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.