Splitting is free

stack and cat allocate, because separate tensors need separate buffers. Splitting is the other direction and costs nothing:

torch.chunk(x, n, dim=0)          # n pieces, all views
torch.split(x, size, dim=0)       # pieces of `size`, all views
x[i * size : (i + 1) * size]      # the same view, written out

Every piece is a change of storage offset and of one axis’s length. Nothing moves.

chunk or split

chunk(x, n) asks for n pieces and works out the sizes, which may be uneven: chunk of a length-7 tensor into 3 gives 3, 3, 1.

split(x, size) asks for pieces of a given size and works out how many, with a short one at the end: split of a length-7 tensor by 3 gives 3, 3, 1 as well, arrived at from the other side.

split also takes a list of sizes, which is the version to know:

q, k, v = qkv.split([d, d, d], dim=-1)

One fused projection, split into three views. No copy, which is why fused QKV is a performance win rather than just tidier.

Micro-batching, and the reason this matters

Splitting a batch to fit in memory is free at the tensor level:

for micro in batch.chunk(4):
    loss = criterion(model(micro), ...) 
    loss.backward()

The pieces are views onto the batch you already have. What costs memory is the activations of each forward pass, not the split, which is precisely why the technique works.

The one caveat

The pieces alias the original, so writing to one writes through, with all the consequences from the indexing track. For reading, which is what a micro-batch does, that is exactly what you want.

Your task

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

Build torch.arange(n), split it into pieces chunks, and return each chunk’s contents plus whether every chunk shares the original’s buffer.