Skip to content

← Shapes and Strides step 8 of 9

Medium Primitives

Sliding windows without copying

Jump to the task

A sliding window selects consecutive elements at regularly spaced start positions. For eight elements, window size three and step two give three complete windows:

input:   [0, 1, 2, 3, 4, 5, 6, 7]
windows: [0, 1, 2]
         [2, 3, 4]
         [4, 5, 6]

A window starting at index six would be incomplete, so it is omitted. The number of windows is (len(x) - size) // step + 1 when 1 <= size <= len(x) and step > 0.

Reuse the storage

Tensor.unfold creates a view instead of copying each window:

import torch

x = torch.arange(8)
y = x.unfold(0, size=3, step=2)
print(list(y.shape))  # [3, 3]
print(y.stride())  # (2, 1)
print(y.tolist())  # [[0, 1, 2], [2, 3, 4], [4, 5, 6]]
print(y.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())  # True

For this contiguous input, y[i, j] reads storage position 2 * i + j. The positions y[0, 2] and y[1, 0] both refer to x[2]. For a strided input whose stride is s, the window strides are (step * s, s).

Stacking slices would allocate number_of_windows * size output elements. The view avoids that data copy, though it still creates metadata. Later operations may allocate outputs or revisit the same input values, so a view does not make the whole calculation cost-free.

Reading and writing overlapping windows

Reductions over these windows are valid. For example, y.sum(dim=1) computes the sum of each window. Summing every entry of y counts repeated input positions once per occurrence; it is not generally the sum of x.

Writes need more care. Changing one entry changes the input storage and any other window entries referring to that location. Avoid bulk in-place updates on overlapping views; clone first if independent writable windows are required. When step >= size, the selected windows do not overlap.

This lesson uses Tensor.unfold, which returns a view. The similarly named torch.nn.Unfold is an image-patch extraction interface and should not be assumed to have the same storage behavior.

Your task

def windows(x: torch.Tensor, size: int, step: int) -> torch.Tensor

For a strided 1-D tensor, return complete windows without copying its storage. Inputs satisfy 1 <= size <= len(x) and step > 0. The output has shape (number_of_windows, size). Drop any trailing incomplete window. The starter gets the values right but stacks them into a separate data buffer.