Transpose a tensor while sharing its storage

This exercise checks two things: the transposed values and whether the result shares the input's storage. A function can return the right values and still fail the storage requirement.

Same values, different storage

import torch

x = torch.arange(6).reshape(2, 3)
view = x.t()
copy = view.contiguous()
print(view.tolist())  # [[0, 3], [1, 4], [2, 5]]
print(torch.equal(view, copy))  # True
print(view.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())  # True
print(copy.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())  # False

For this input, t() returns a view. It swaps the two dimensions and their strides without copying values. The following contiguous() call creates a separate data buffer because the transpose is not contiguous.

A copy uses extra memory and takes time. Whether that matters depends on the tensor size and workload. This exercise asks you to avoid the copy because the transpose can already be represented as a view.

Read the layout

A stride tells you how many storage elements to advance when an axis index increases by one. Here x has shape (2, 3) and strides (3, 1). Its transpose has shape (3, 2) and strides (1, 3). Both descriptions refer to the same six stored values.

The storage offset also matters: a slice may begin partway into a buffer. For the live, nonempty tensors on the same device used here, comparing storage pointers checks whether the buffer is shared. Comparing tensor data_ptr() values instead checks first-element addresses, which can differ between views of the same storage.

Contiguous does not always mean copied

If a tensor is already contiguous in the requested memory format, .contiguous() returns it unchanged. Transposing a single-row or single-column matrix can leave it contiguous, so the starter passes those cases. It still copies the other test inputs.

def transpose(x: torch.Tensor) -> torch.Tensor