Add row sums without changing the destination shape

An ordinary addition can produce a broadcasted result with a larger shape. An in-place addition must keep the destination's existing shape.

import torch

a = torch.zeros(3, 4)
b = torch.ones(3, 1)
print((a + b).shape)  # torch.Size([3, 4])
a.add_(b)            # Works: b broadcasts to a's shape.

Reversing the in-place operation does not work. b.add_(a) would need a (3, 4) result, but b has shape (3, 1). PyTorch raises a RuntimeError instead of resizing b.

Broadcasting treats compatible singleton dimensions of the source as repeated values without materializing repeated data. The destination of add_() must already have the result's shape. Shape compatibility is only one requirement: dtype and autograd rules also apply to in-place operations.

Accumulating one total per row

Suppose totals has shape (rows, 1) and batch has shape (rows, cols). You want to add each row's sum to its existing total. Adding the whole batch is not the same calculation, and cannot fit when cols is greater than one. Reduce each row first:

totals = torch.tensor([[10.0], [20.0]])
batch = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
row_sums = batch.sum(dim=1, keepdim=True)
print(row_sums.tolist())  # [[3.0], [7.0]]
totals.add_(row_sums)
print(totals.tolist())    # [[13.0], [27.0]]

keepdim=True retains the reduced axis, producing (rows, 1) rather than (rows,). That matches totals directly and avoids accidentally aligning the row count with its last axis.

Match shapes from the right

Broadcasting lets an operation apply across compatible shapes without explicitly constructing repeated inputs. Compare dimensions from the right: each pair must be equal, or one must be 1. Missing leading dimensions behave like dimensions of 1. At each position, equal sizes stay the same. If one size is 1, the result uses the other size. This also covers empty dimensions: (0, 3) and (1, 3) broadcast to (0, 3), not (1, 3).

For example, (4, 3) plus (3,) applies the same three-feature bias to all four rows. Adding (4, 1) instead applies a different scalar to each row. A size-one axis lets the same value be reused along that axis.

This explains a subtle centring bug: reducing (B, D) over features gives (B,). Subtracting it directly aligns B with the feature axis. If B == D, the operation runs but uses the wrong means. Retaining the reduction axis as (B, 1) expresses the intended dependence.

Check how much memory the result needs

The input need not be physically copied, but the output can still be huge: (N, 1, D) minus (1, M, D) has shape (N, M, D). An in-place operation must fit its result in the existing destination shape. Explicit expansion may also create a view whose positions refer to the same storage, making independent writes unsafe.

Before implementing an operation, write the aligned shapes and the result shape. Then check which inputs are reused and how much space the output needs.