Foundations

Masked Log-Sum-Exp

EasyRaw PyTorchAbout 10 min1 public example plus private tests

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

Your task

For each slice along dim, return the natural logarithm of the sum of exponentials of its included scores. This is the log normalizing constant used by softmax; dim is reduced.

Function signature

masked_logsumexp(
    logits: torch.Tensor,
    mask: torch.Tensor,
    dim: int = -1,
    sentinel: float = 0.0,
    keepdim: bool = False,
) -> torch.Tensor

Worked example

Unequal included scores

Take logits=[0,1,100], mask=[True,True,False] and dim=0. Only exponential weights 1 and e participate; the masked score 100 is excluded even from the denominator. The result is log(1+e) ≈ 1.313262. With keepdim=False it is a scalar; with keepdim=True it is a one-element tensor. This is not the logarithm of the score sum.

Parameters

Parameter	Shape and meaning
`logits`torch.Tensor	$[\mathrm{d0}, \ldots, d(R-1)]$ Finite values whose exponentials are summed; the reduction axis may be empty.
`mask`torch.Tensor	Any shape broadcastable to logits.shape.True includes an entry in the exponential sum; False excludes it.
`dim`intDefault `-1`	Axis to reduce; negative indices count from the end.
`sentinel`floatDefault `0.0`	Finite return value for an entirely masked or empty reduction slice.
`keepdim`boolDefault `False`	True retains the reduced axis with size one; False removes it.

Return

torch.Tensor · logits.shape with dim removed, or retained at size 1 when keepdim=True.

One log exponential sum per slice along dim. Remove dim unless keepdim=True, which retains it at size one. Entirely masked or zero-length reduction slices return sentinel.

Required behavior

logits has rank at least one; mask is Boolean and must broadcast to exactly logits.shape. True includes a score. dim is one integer axis (default -1); negative axes count from the end. A slice fixes every index except dim.
For a nonempty included set V, let a be its largest score and Z = sum(exp(logits[i]-a) for i in V). Subtract this included maximum before exponentiating; masked scores contribute neither exponential weight nor normalizer.
Return a+log(Z) for each nonempty included slice. Entirely masked or zero-length slices return sentinel (default 0) at every retained output index. Remove dim unless keepdim=True, which retains it with size one.
Use no epsilon, probability clipping or probability floor. Empty-support outputs have zero gradients. Preserve the shared CPU float32/float64 dtype and device, accept noncontiguous inputs, and do not mutate inputs.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

Why this matters

Subtracting a valid maximum controls exponent size; an explicit empty-row branch prevents an all-masked input from becoming NaN or a uniform distribution.

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_logsumexp(
    logits: torch.Tensor,  # (d0, ..., d(R-1))
    mask: torch.Tensor,  # Any shape broadcastable to logits.shape.
    dim: int = -1,
    sentinel: float = 0.0,
    keepdim: bool = False,
) -> torch.Tensor:
    """Return a tensor with shape logits.shape with dim removed, or retained at size 1 when
        keepdim=True."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement masked_logsumexp")

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

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