Skip to content

← Memory, Copies and In-Place step 2 of 6

Easy Primitives

Write into the buffer you were given

Most PyTorch functions allocate their result. Many of them will instead write into a buffer you supply:

torch.add(a, b)              # allocates
torch.add(a, b, out=dest)    # writes into dest, returns dest

The second dispatches the same single kernel. What it saves is the allocation, and in a loop that runs every step, the allocation is the part that shows up.

Why not just use the in-place form

a.add_(b)       # writes into a
torch.add(a, b, out=dest)   # writes into dest, leaves a and b alone

add_ destroys an operand. out= writes somewhere else entirely, which is what you want when both inputs are still needed, or when the destination is a slice of a larger buffer you are filling.

Where it earns its keep

A preallocated output buffer that gets refilled every iteration:

buf = torch.empty(batch, dim)
for batch_in in loader:
    torch.matmul(batch_in, w, out=buf)
    ...

One allocation for the whole loop instead of one per step. This is the shape of most inference servers and most custom training loops.

The caveats

out= requires the destination to have the right shape and dtype. It will resize it with a warning if not, which is a sign your buffer is not doing the job you allocated it for.

It is also not magic: PyTorch’s allocator caches, so a freshly allocated tensor in a loop is usually cheap. out= matters most when the tensors are large, when the loop is tight, or when you want a stable address to hand to something outside PyTorch.

Your task

def scale_into(dest: torch.Tensor, x: torch.Tensor, factor: float) -> torch.Tensor

Write x * factor into dest and return dest. dest must still be the same buffer afterwards, and x must be unchanged.