Foundations

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, add the included values and divide by their count plus eps. Return zero when that count is zero.

Function signature

masked_mean(
    x: torch.Tensor,
    mask: 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 includes 2 and 4. With eps=0 in this example, its mean is (2+4)/2 = 3. 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)
dim = -1
keepdim = False
eps = 0.0

result = masked_mean(x=x, mask=mask, dim=dim, keepdim=keepdim, eps=eps)

Expected return

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

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[d_0,\ldots,d_{R-1}]$ with arbitrary rank $R\ge1$ .Values to average.
`mask`torch.Tensor	Any shape broadcastable to `x.shape`.True includes an entry; False excludes it.
`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 offset added once to the included-entry count. Empty slices return 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 Mean of the included values, with the requested axes reduced.

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.
Let $S$ be the sum of included values and $n$ their count after broadcasting. Return $S/(n+\varepsilon)$ when $n>0$ , and zero otherwise.
eps is finite and nonnegative, default $10^{-8}$ . Add it once to the count after reduction. Do not average smaller means or count the unexpanded mask.
Entirely masked slices have zero gradients. Make inactive divisions safe as well as the selected result.

Why this matters

Masking lets padded batches share one tensor while averages use only real entries.

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_mean(
    x: torch.Tensor,  # (d_0,...,d_{R-1}) with arbitrary rank R>=1.
    mask: 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 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.