Foundations

Masked Sum

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

Sum the included entries of x over the requested axes. Return a tensor containing one sum for each remaining index.

Function signature

masked_sum(
    x: torch.Tensor,
    mask: torch.Tensor,
    dim: int | list[int] | tuple[int, ...] | None = -1,
    keepdim: bool = False,
) -> torch.Tensor

Worked example

One included row and one empty row

The first row contributes $2+4=6$ . No entries participate in the second row, so its sum is 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

result = masked_sum(x=x, mask=mask, dim=dim, keepdim=keepdim)

Expected return

torch.tensor([6.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 add at included positions.
`mask`torch.Tensor	Any shape broadcastable to `x.shape`.True includes an entry in the sum; False excludes it.
`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.

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 Sum of included values over the requested axes.

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.
Preserve gradients for included values; excluded values have zero derivative.

Why this matters

A mask lets padded entries share a tensor with real data without contributing to the sum.

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_sum(
    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,
) -> 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_sum")

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.