Tensor Reshaping

Tensors are the basic data containers used in PyTorch, and most of the work done with them involves changing how their values are arranged. This article follows one short program in which a single small tensor is created, inspected, reshaped, transposed, and multiplied by its own transpose. Each operation is simple on its own. The reason to study them together is that one wrong assumption early in the program, such as a shape you did not expect, carries forward into every later line and can look like a failure somewhere else entirely.

The starting tensor is built from a list of two inner lists, each holding three numbers. PyTorch reads that nested structure and creates a two-dimensional tensor. Its shape attribute reports how many entries lie along each dimension, which is the first thing to check whenever a later operation behaves strangely. The program below runs the full sequence. As you read it, predict the printed result of each line before you look at the output.

The reshape method takes a new shape and returns a tensor with that shape. The starting tensor holds six values, so any new shape must also account for exactly six entries. Asking for a shape with a different total count cannot succeed. Reshaping does not shuffle or recalculate anything. The six values keep their sequence, and the new shape decides where each row begins and ends. Reading the same six values in order into three rows of two produces a different grouping of the same numbers, as the first grid below shows.

Under the grids, the values are stored as one linear sequence, and the shape only tells the library how to index into that sequence. Reshape keeps that sequence intact and changes the indexing rule, which is why the values stay the same. The source's reshaped result is three rows of two: the pairs 1 and 2, 3 and 4, and 5 and 6. Those are the same six numbers in the same order as the original, regrouped.

The transpose attribute, written .T with no parentheses, flips a two-dimensional tensor so that its rows become its columns and its columns become its rows. The first row of the original becomes the first column of the result, and the same pattern continues for every row. Transposing changes the shape from 2 by 3 to 3 by 2. In the diagram below, each element is labeled with its new position, and the elements that moved are marked.

The matmul method multiplies two tensors as matrices. In the program above, the tensor is multiplied by its own transpose. The original is 2 by 3 and the transpose is 3 by 2. The inner dimensions, the 3 at the end of the first shape and the 3 at the start of the second, agree, so the product is defined. The outer dimensions, the 2 and the 2, give the shape of the result. Before running matmul, check that the inner dimensions match. If they do not, the operation cannot succeed.

Each output entry combines one row of the first tensor with one column of the second. The source prints the product with rows 14 and 32 in the first row and 32 and 77 in the second. To check the top-left entry by hand, multiply the first row by the first column and add the products: 1 times 1, plus 2 times 2, plus 3 times 3, which equals 14. Checking a single entry this way is a fast test when a result looks wrong, because it takes only a few multiplications.

Shape problems are the most common source of confusing tensor results. The mistakes below are easy to make because each one looks reasonable until the printed shape or values are compared with what was expected.

When an output differs from what you expected, the efficient approach is to move forward through the program one operation at a time. Predict the shape and values after each line, compare them with what is printed, and stop at the first mismatch. That line is the failure location. Later lines may be consequences of it rather than separate causes.