Skip to content

← Batching and the Data Path step 4 of 5

Easy Primitives

Split without copying

Jump to the task

Splitting a tensor with torch.chunk or torch.split returns views that share its storage. PyTorch creates tensor objects describing the pieces, but does not copy their elements.

import torch

x = torch.arange(7)
chunks = torch.chunk(x, 3)
print([part.tolist() for part in chunks])
# [[0, 1, 2], [3, 4, 5], [6]]

Each view records its shape, strides and starting offset into the shared storage. Splitting has bookkeeping overhead, so “no data copy” does not mean “no cost.” Joining tensors with ordinary cat or stack generally creates an output containing copied data.

Requested chunks versus chunk size

torch.chunk(x, chunks) requests a number of pieces, but may return fewer. For example:

parts = torch.chunk(torch.arange(6), 4)
print([part.tolist() for part in parts])
# [[0, 1], [2, 3], [4, 5]]

torch.split(x, split_size) instead sets the maximum size of each piece. The last piece can be shorter. It also accepts a list of exact sizes that must add up to the length of the selected dimension.

Splitting a combined query/key/value projection this way avoids copying its output into three new tensors. Whether a fused projection is faster depends on the model, tensor sizes and backend.

Shared storage still matters

Changing a piece changes the corresponding elements of the original tensor:

chunks[0].add_(10)
print(x.tolist())  # [10, 11, 12, 3, 4, 5, 6]

Views can help process a batch in smaller pieces without copying the input. The full input storage remains alive while its views are used. Memory savings come from limiting intermediate activations, provided you release each piece's autograd graph instead of retaining all of them. Matching a full-batch gradient also requires appropriate loss weighting and compatible model behavior.

Your task

def split_batch(n: int, pieces: int) -> dict

For positive integers n and pieces, build torch.arange(n, dtype=torch.float32) and split it with torch.chunk. Return a dictionary containing:

  • "chunks": each returned chunk converted to a Python list.
  • "all_share": whether every chunk shares the original tensor's storage.

Return the chunks PyTorch actually produces, even if there are fewer than pieces. Compare storage addresses, since a later chunk's first element starts at a different offset within the same storage.