Shaping Up: A Hands-On Guide to Tensor Manipulation in PyTorch

Learn how tensors work in PyTorch, from simple scalars to 3D arrays. This notebook walks through creating and reshaping tensors, understanding dimensions, and using operations like squeeze, unsqueeze, ravel, reshape, transpose, and stacking/splitting to manipulate tensor shapes.

pytorch
neural-networks
Author

Mei-Chin Pang

Published

March 7, 2026

This notebook provides a hands-on introduction to tensor manipulation in PyTorch and the fundamental data structure used throughout the framework.

What Are Tensors?

A tensor is the fundamental data structure in PyTorch, which is a generalisation of scalars, vectors, and matrices to arbitrary dimensions.

Each step up in dimensionality nests one more level of structure:

  • A scalar is a single value.
  • A vector is a list of scalars.
  • A matrix is a list of vectors (rows).
  • A tensor is a list of matrices (or more generally, a list of lower-dimensional tensors).

In PyTorch, all of these are represented by torch.Tensor.

Type Dimensions Shape example Description
Scalar 0 () A single number, e.g. 11
Vector 1 (3,) A 1D array of numbers, e.g. [5, 3, 7]. Can be a row vector (shape 1×3) or column vector (shape 3×1)
Matrix 2 (2, 3) A 2D grid of numbers with rows and columns, e.g. a 2×3 table ( 2 rows, 3 columns)
Tensor 3+ (2, 3, 3) A higher-dimensional array, we can think of it as a stack of matrices, or a cube of numbers

What Are Tensors (Source: Tensors 101: The Building Blocks of PyTorch)
import torch
import numpy as np

Example: Creating Scalars, Vectors, Matrices, and Tensors: NumPy vs PyTorch

Type Matrix Notation NumPy PyTorch Dimensions
Scalar
x \in \mathbb{R}		 x = 11 np.array(11) torch.tensor(11) 0D: shape (). A single number, no axes.
Vector
\mathbf{v} \in \mathbb{R}^{4}		 \mathbf{v} = \begin{bmatrix} 5 & 3 & 7 & 1 \end{bmatrix} np.array([5, 3, 7, 1]) torch.tensor([5, 3, 7, 1]) 1D: shape (4,). One axis (dim 0) with 4 elements.
Row vector
\mathbf{v}^T \in \mathbb{R}^{1 \times 4}		 \mathbf{v}^T = \begin{bmatrix} 5 & 3 & 7 & 1 \end{bmatrix} np.array([[5, 3, 7, 1]]) torch.tensor([[5, 3, 7, 1]]) 2D: shape (1, 4).
Dim 0: 1 row.
Dim 1: 4 columns.
Column vector
\mathbf{v} \in \mathbb{R}^{4 \times 1}		 \mathbf{v} = \begin{bmatrix} 5 \\ 3 \\ 7 \\ 1 \end{bmatrix} np.array([[5], [3], [7], [1]]) torch.tensor([[5], [3], [7], [1]]) 2D: shape (4, 1).
Dim 0: 4 rows
Dim 1: 1 column.
Matrix
\mathbf{M} \in \mathbb{R}^{2 \times 4}		 \mathbf{M} = \begin{bmatrix} 1 & 2 & 3 & 4 \\ 5 & 6 & 7 & 8 \end{bmatrix} np.array([[1,2,3,4], [5,6,7,8]]) torch.tensor([[1,2,3,4], [5,6,7,8]]) 2D: shape (2, 4).
Dim 0: 2 rows
Dim 1: 4 columns.
3D Tensor
\mathcal{T} \in \mathbb{R}^{2 \times 3 \times 4}		 \mathcal{T}_{[0]} = \begin{bmatrix} 1 & 2 & 3 & 4 \\ 5 & 6 & 7 & 8 \\ 9 & 10 & 11 & 12 \end{bmatrix}, \mathcal{T}_{[1]} = \begin{bmatrix} 13 & 14 & 15 & 16 \\ 17 & 18 & 19 & 20 \\ 21 & 22 & 23 & 24 \end{bmatrix} np.array([[[1,2,3,4], ...]]) torch.tensor([[[1,2,3,4], ...]]) 3D: shape (2, 3, 4).
Dim 0: 2 blocks
Dim 1: 3 rows per block
Dim 2: 4 columns per row.

Key differences:

Feature NumPy PyTorch
Constructor np.array(...) torch.tensor(...)
Shape attribute .shape .shape or .size()
Dimensions count .ndim .dim()
GPU support No Yes 
# Numpy Scalar (0D array)
np_scalar = np.array(11)
print(f"Numpy Scalar: {np_scalar}")
print(f"shape: {np_scalar.shape}, dim: {np_scalar.ndim}")
print("-"*79)

# One-dimensional array with four elements
# Create with a list of elements
np_vector = np.array([5, 3, 7, 1])
print(f"Numpy Vector: {np_vector}")
print(f"shape: {np_vector.shape}, dim: {np_vector.ndim}")
print("-"*79)

# Two-dimensional array (row vector) with 1 row and 4 columns
# Only one row, so the inner list represents that single row
np_row_vector = np.array([[5, 3, 7, 1]])
print(f"Numpy Row Vector:\n{np_row_vector}")
print(f"shape: {np_row_vector.shape}, dim: {np_row_vector.ndim}")
print("-"*79)

# Two-dimensional array (column vector) with 4 rows and 1 column
# Only one column, so each inner list represents a single 
# element in that column
np_vector_column = np.array([[5], [3], [7], [1]])
print(f"Numpy Column Vector:\n{np_vector_column}")
print(f"shape: {np_vector_column.shape}, dim: {np_vector_column.ndim}")
print("-"*79)

# Two-dimensional array (matrix) with 2 rows and 4 columns
# Two rows, so each inner list represents a row
np_matrix = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
print(f"Numpy Matrix:\n{np_matrix}")
print(f"shape: {np_matrix.shape}, dim: {np_matrix.ndim}")
print("-"*79)

# Three-dimensional array (tensor) with 2 blocks, 
# each containing 3 rows and 4 columns
np_tensor = np.array([
    [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
    [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]])
print(f"Numpy Tensor:\n{np_tensor}")
print(f"shape: {np_tensor.shape}, dim: {np_tensor.ndim}")
print("-"*79)
Numpy Scalar: 11
shape: (), dim: 0
-------------------------------------------------------------------------------
Numpy Vector: [5 3 7 1]
shape: (4,), dim: 1
-------------------------------------------------------------------------------
Numpy Row Vector:
[[5 3 7 1]]
shape: (1, 4), dim: 2
-------------------------------------------------------------------------------
Numpy Column Vector:
[[5]
 [3]
 [7]
 [1]]
shape: (4, 1), dim: 2
-------------------------------------------------------------------------------
Numpy Matrix:
[[1 2 3 4]
 [5 6 7 8]]
shape: (2, 4), dim: 2
-------------------------------------------------------------------------------
Numpy Tensor:
[[[ 1  2  3  4]
  [ 5  6  7  8]
  [ 9 10 11 12]]

 [[13 14 15 16]
  [17 18 19 20]
  [21 22 23 24]]]
shape: (2, 3, 4), dim: 3
-------------------------------------------------------------------------------
# Scalar (0D tensor)
torch_scalar = torch.tensor(11)
print(f"Torch scalar: {torch_scalar}") 
print(f"shape: {torch_scalar.shape}, dim: {torch_scalar.dim()}")
print("-"*79)

# Vector (1D tensor)
# Create with a list of elements
torch_vector = torch.tensor([5, 3, 7, 1])
print(f"Torch vector: {torch_vector}")
print(f"shape: {torch_vector.shape}, dim: {torch_vector.dim()}")
print("-"*79)

# Row vector (2D tensor) with 1 row and 4 columns
# Only one row, so the inner list represents that single row
torch_row_vector = torch.tensor([[5, 3, 7, 1]])
print(f"Torch row vector:\n{torch_row_vector}")
print(f"shape: {torch_row_vector.shape}, dim: {torch_row_vector.dim()}")
print("-"*79)

# Two-dimensional array (column vector) with 4 rows and 1 column
# Only one column, so each inner list represents a single 
# element in that column
torch_column_vector = torch.tensor([[5], [3], [7], [1]])
print(f"Torch column vector:\n{torch_column_vector}")
print(f"shape: {torch_column_vector.shape}, dim: {torch_column_vector.dim()}")
print("-"*79)

# Matrix (2D tensor)
# Two rows, so each inner list represents a row
torch_matrix = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]])
print(f"Torch matrix:\n{torch_matrix}")
print(f"shape: {torch_matrix.shape}, dim: {torch_matrix.dim()}")
print("-"*79)

# Tensor (3D tensor) with 2 blocks, each containing 3 rows and 4 columns
# Outer list represents the blocks
# Second outer list represents the rows within each block
# Innermost list represents the columns within each row
torch_tensor_3d = torch.tensor([
    [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
    [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]])
print(f"Torch 3D tensor:\n{torch_tensor_3d}")
print(f"shape: {torch_tensor_3d.shape}, dim: {torch_tensor_3d.dim()}")
Torch scalar: 11
shape: torch.Size([]), dim: 0
-------------------------------------------------------------------------------
Torch vector: tensor([5, 3, 7, 1])
shape: torch.Size([4]), dim: 1
-------------------------------------------------------------------------------
Torch row vector:
tensor([[5, 3, 7, 1]])
shape: torch.Size([1, 4]), dim: 2
-------------------------------------------------------------------------------
Torch column vector:
tensor([[5],
        [3],
        [7],
        [1]])
shape: torch.Size([4, 1]), dim: 2
-------------------------------------------------------------------------------
Torch matrix:
tensor([[1, 2, 3, 4],
        [5, 6, 7, 8]])
shape: torch.Size([2, 4]), dim: 2
-------------------------------------------------------------------------------
Torch 3D tensor:
tensor([[[ 1,  2,  3,  4],
         [ 5,  6,  7,  8],
         [ 9, 10, 11, 12]],

        [[13, 14, 15, 16],
         [17, 18, 19, 20],
         [21, 22, 23, 24]]])
shape: torch.Size([2, 3, 4]), dim: 3

Understanding Tensor Dimensions: How to Read the Dimensions

A tensor’s shape describes its size along each dimension (also called an axis). Understanding the tensor’s dimensions is not necessarily straight-forward, the following two examples illustrate how to read tensor of different dimensions.

An Illustration of Tensor Dimension Example (Source: Everything You Need to Know About Tensors in PyTorch)

Example (1): Consider the tensor with torch.Size([1, 3, 3])

tensor([[[1, 2, 3],
         [3, 6, 9],
         [2, 4, 5]]])

torch.Size([1, 3, 3])

For a tensor with shape torch.Size([1, 3, 3]), there are three dimensions:

Dimension (dim) Index Size What it represents
0 Outermost 1 Number of “blocks” (e.g. 1 block).
1 Middle 3 Rows within each block: here there are 3 rows: [1,2,3], [3,6,9], [2,4,5].
2 Innermost 3 Elements within each row: here there are 3 columns per row.

This tensor contains 1 \times 3 \times 3 = 9 elements in total.

Example (2): Consider the tensor with the torch.Size([3, 4, 5])

tensor([[[-1.2748,  0.0892,  1.7380,  2.1273, -1.6192],
         [-1.0360,  2.1186,  0.4506,  1.4779, -0.3235],
         [-0.8709,  1.3224,  0.1082,  1.1610, -0.5307],
         [ 0.1087, -1.7062,  0.9570, -0.5008,  0.9883]],

        [[-0.9094, -0.6152,  1.7460,  0.9939, -0.0149],
         [-0.3225, -1.8547,  0.0158, -0.8520, -0.3061],
         [ 0.1341, -0.1355, -1.4467, -0.0038,  1.0789],
         [-1.4798, -0.7031, -0.1879, -0.9284,  0.4436]],

        [[ 0.5743,  1.7551,  0.1351,  0.1438,  0.7122],
         [-0.0231, -0.0894, -0.2458, -0.1292,  0.4462],
         [ 0.3520, -1.1397,  1.3075, -1.0896,  0.1002],
         [ 0.1951, -0.2193,  0.4690, -1.4147,  0.1399]]])
torch.Size([3, 4, 5])

This tensor created by torch.randn(3, 4, 5) has shape (3, 4, 5), which has again three dimensions, but with different sizes:

Dimension (dim) Index Size What it represents
0 Outermost 3 3 blocks (e.g. 3 matrices)
1 Middle 4 4 rows within each block
2 Innermost 5 5 columns (elements per row)

This tensor contains 3 \times 4 \times 5 = 60 elements in total. We can think of it as 3 matrices, each with 4 rows and 5 columns.

General Rule:

The number of dimensions equals the nesting depth of brackets. Each dimension is numbered from the outermost (0) to the innermost. The .shape attribute (or .size()) returns the size along each dimension as a tuple.

squeeze() and unsqueeze(): Adding and Removing Size-1 Dimensions

unsqueeze(dim) and squeeze() are inverse operations that add or remove dimensions of size 1 without changing the underlying data, which are commonly used to match tensor ranks for broadcasting, batching, or matrix multiplication. Both operations return a view (shared memory, no data copy). squeeze(dim) can also target a specific dimension, it only removes it if its size is 1.

An Illustration of Squeeze vs Unsqueeze Operation (Source: squeeze() vs unsqueeze() in PyTorch)

unsqueeze(dim): add a dimension

Inserts a new dimension of size 1 at the specified position. Starting from a 2D tensor with shape (2, 2):

Operation Resulting shape Effect
unsqueeze(0) (1, 2, 2) Adds a dimension at the front by wrapping the matrix in an outer bracket
unsqueeze(1) (2, 1, 2) Adds a dimension in the middle, where each row becomes its own “block”
unsqueeze(2) (2, 2, 1) Adds a dimension at the end, where each element becomes a single-element list

In all three cases, the data ([[1, 2], [3, 4]]) is unchanged. Only the shape metadata is modified, turning a 2D tensor into a 3D tensor.

squeeze(): remove dimensions

Removes all dimensions of size 1. Applying squeeze() to any of the three 3D results above returns the original (2, 2) matrix, since each has exactly one size-1 dimension.

# 2D matrix from the image
a = torch.tensor([[1, 2], [3, 4]])
print(f"Original: {a.shape}")
print(a)
print("-"*79)

# unsqueeze at dim=0:(1, 2, 2)
u0 = a.unsqueeze(0)
print(f"unsqueeze(0): {u0.shape}")
print(u0)
print("-"*79)

# unsqueeze at dim=1:(2, 1, 2)
u1 = a.unsqueeze(1)
print(f"unsqueeze(1): {u1.shape}")
print(u1)
print("-"*79)

# unsqueeze at dim=2:(2, 2, 1)
u2 = a.unsqueeze(2)
print(f"unsqueeze(2): {u2.shape}")
print(u2)
print("-"*79)

# squeeze reverses each unsqueeze back to (2, 2)
print(f"squeeze(u0): {u0.squeeze().shape}")
print(f"squeeze(u1): {u1.squeeze().shape}")
print(f"squeeze(u2): {u2.squeeze().shape}")
Original: torch.Size([2, 2])
tensor([[1, 2],
        [3, 4]])
-------------------------------------------------------------------------------
unsqueeze(0): torch.Size([1, 2, 2])
tensor([[[1, 2],
         [3, 4]]])
-------------------------------------------------------------------------------
unsqueeze(1): torch.Size([2, 1, 2])
tensor([[[1, 2]],

        [[3, 4]]])
-------------------------------------------------------------------------------
unsqueeze(2): torch.Size([2, 2, 1])
tensor([[[1],
         [2]],

        [[3],
         [4]]])
-------------------------------------------------------------------------------
squeeze(u0): torch.Size([2, 2])
squeeze(u1): torch.Size([2, 2])
squeeze(u2): torch.Size([2, 2])

Advanced Boolean Indexing with PyTorch

This example demonstrates advanced boolean indexing to select columns or rows based on a condition applied across an axis.

Boolean Column Filtering

Breakdown of a[:, (a > 0).all(axis=0)]:

Expression Description
a > 0 Creates a boolean tensor of the same shape as a, where each element is True if it is greater than 0
.all(axis=0) A value is True only if all elements in that column satisfy the condition
a[:, ...] Selects all rows, but only the columns where the condition is True

Replace .all() with .any() to select columns where at least one value satisfies the condition.

Boolean Row Filtering

Breakdown of a[(a > 0).all(axis=1)]:

Expression Description
a > 0 Creates a boolean tensor of the same shape as a, where each element is True if it is greater than 0
.all(axis=1) A value is True only if all elements in that row satisfy the condition
a[...] Selects only the rows where the condition is True

Replace .all() with .any() to select rows where at least one value satisfies the condition.

# Example of boolean indexing
a = torch.tensor([
    [ 0.6351,  0.6529, -0.6906, -0.0558],
    [ 0.4105, -0.3270, -0.4520,  0.5239],
    [ 0.4299,  0.9722, -0.1614,  0.5081]])
print(a)
print("-"*79)

# Column filtering
print("A column is selected if all its values are > 0:")
print(a[:, (a > 0).all(axis=0)])
print("-"*79)
print("A column is selected if any of its values are > 0:")
print(a[:, (a > 0).any(axis=0)])
print("-"*79)

# Row filtering
print("A row is selected if all its values are > 0:")
print(a[(a > 0).all(axis=1)])
print("-"*79)
print("A row is selected if any of its values are > 0:")
print(a[(a > 0).any(axis=1)])
tensor([[ 0.6351,  0.6529, -0.6906, -0.0558],
        [ 0.4105, -0.3270, -0.4520,  0.5239],
        [ 0.4299,  0.9722, -0.1614,  0.5081]])
-------------------------------------------------------------------------------
A column is selected if all its values are > 0:
tensor([[0.6351],
        [0.4105],
        [0.4299]])
-------------------------------------------------------------------------------
A column is selected if any of its values are > 0:
tensor([[ 0.6351,  0.6529, -0.0558],
        [ 0.4105, -0.3270,  0.5239],
        [ 0.4299,  0.9722,  0.5081]])
-------------------------------------------------------------------------------
A row is selected if all its values are > 0:
tensor([], size=(0, 4))
-------------------------------------------------------------------------------
A row is selected if any of its values are > 0:
tensor([[ 0.6351,  0.6529, -0.6906, -0.0558],
        [ 0.4105, -0.3270, -0.4520,  0.5239],
        [ 0.4299,  0.9722, -0.1614,  0.5081]])

Flattening Tensors into a Single Dimension with Tensor.ravel()

  • a.ravel() returns a contiguous 1D flattened tensor containing all elements of a. It is equivalent to a.reshape(-1).
  • For a tensor with shape (3, 4), .ravel() produces a 1D tensor of shape (12,).
  • Elements are read in row-major (C) order.
  • If the tensor is already contiguous in memory, .ravel() returns a view (no data copy). Otherwise, it returns a copy.

What does “contiguous” mean?

A tensor is contiguous when its elements are stored in a single, unbroken block of memory in the same order we visit them iterating dimensions from first to last (row-major/C order).

Many operations (view, ravel) just reinterpret existing memory with a new shape. This only works if elements are already contiguous. If not, PyTorch must copy the data first.

Operations that change logical order without copying data can break contiguity:

Contiguous Non-contiguous
Memory layout Elements in sequential order Elements scattered or reordered
view()
ravel()		 Returns a view (fast, no copy) Raises an error or must copy
Example cause Default creation, .contiguous() .t(), .permute(), .narrow()

We can check with a.is_contiguous() and force a contiguous copy with a.contiguous().

What does Row-major C mean?

Here, Row-major (C) order describes how a multi-dimensional array is laid out as a flat sequence in memory. Each row is stored contiguously, and rows are placed one after another, hence “row-major.”

In the context of ravel(), when we call a.ravel() on the tensor above, elements are read row by row:

[a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23]

This matches exactly how the data already sits in memory, which is why ravel() can return a view (no copy needed) when the tensor is contiguous.

a = torch.randn(3,4)
print("Original tensor:")
print(a)
print("-"*79)

print("Raveled tensor:")
b = a.ravel()
print(b)
print(f"Is b contiguous? {b.is_contiguous()}")
print("-"*79)

print("Reshaped tensor:")
c = a.reshape(-1)
print(c)
print(f"Is c contiguous? {c.is_contiguous()}")
Original tensor:
tensor([[-0.1289,  0.7486, -0.9761, -0.0435],
        [-0.6773,  0.9603, -0.5627,  0.1547],
        [ 0.9936, -0.5449,  1.8688,  0.9037]])
-------------------------------------------------------------------------------
Raveled tensor:
tensor([-0.1289,  0.7486, -0.9761, -0.0435, -0.6773,  0.9603, -0.5627,  0.1547,
         0.9936, -0.5449,  1.8688,  0.9037])
Is b contiguous? True
-------------------------------------------------------------------------------
Reshaped tensor:
tensor([-0.1289,  0.7486, -0.9761, -0.0435, -0.6773,  0.9603, -0.5627,  0.1547,
         0.9936, -0.5449,  1.8688,  0.9037])
Is c contiguous? True

Transpose a 2D tensor will change the memory layout, so the result after transpose is not contiguous.

a = torch.randn(3, 4)
print(f"Is a contiguous? {a.is_contiguous()}")

b = a.t()
print(f"Is b contiguous? {b.is_contiguous()}")
Is a contiguous? True
Is b contiguous? False

Reshaping Tensors with Tensor.reshape()

  • a.reshape(*shape) returns a tensor with the same data but a new shape. The total number of elements must remain the same.
  • If the original tensor is contiguous, reshape returns a view (no copy, shares memory).
  • If not contiguous, it returns a copy with the new shape.
  • We can use -1 for one dimension to let PyTorch infer it automatically: a.reshape(3, -1, 2) gives (3, 2, 2).

Example:

a = torch.randn(3, 4)
a.reshape(3, 2, 2)

The original tensor a has a shape of (3, 4) with 12 elements. reshape(3, 2, 2) rearranges it into shape (3, 2, 2), which still retains the 12 elements (3 \times 2 \times 2 = 12).

a = torch.randn(3,4)
print("Original tensor:")
print(a)
print(a.shape)
print("-"*79)

print("Reshaped tensor (3,2,2):")
print(a.reshape(3,2,2))
print(a.reshape(3,2,2).shape)
print("-"*79)

print("Reshaped tensor (-1), flatten the tensor to 1D:")
print(a.reshape(-1))
print(a.reshape(-1).shape)
Original tensor:
tensor([[-1.0349, -0.7509,  0.2990, -0.3119],
        [-0.2913,  0.0147,  0.7607,  0.7079],
        [-0.0297, -0.7357, -0.7373,  0.1570]])
torch.Size([3, 4])
-------------------------------------------------------------------------------
Reshaped tensor (3,2,2):
tensor([[[-1.0349, -0.7509],
         [ 0.2990, -0.3119]],

        [[-0.2913,  0.0147],
         [ 0.7607,  0.7079]],

        [[-0.0297, -0.7357],
         [-0.7373,  0.1570]]])
torch.Size([3, 2, 2])
-------------------------------------------------------------------------------
Reshaped tensor (-1), flatten the tensor to 1D:
tensor([-1.0349, -0.7509,  0.2990, -0.3119, -0.2913,  0.0147,  0.7607,  0.7079,
        -0.0297, -0.7357, -0.7373,  0.1570])
torch.Size([12])

Transposing a Matrix with Tensor.T and torch.transpose()

PyTorch provides two common ways to transpose a 2D tensor (matrix):

  • Tensor.T
  • torch.transpose(input, dim0, dim1)

Both operations return a view of the original tensor. They share the same underlying data, so modifying one affects the other. The resulting transposed tensor is not contiguous in memory. Call .contiguous() if a subsequent operation requires it.

Tensor.T

A property that reverses all dimensions of the tensor, where rows become columns and vice versa.

a.T  # shape (3, 4) → (4, 3)

torch.transpose(input, dim0, dim1)

Swaps two specific dimensions of a tensor. For a 2D matrix, torch.transpose(a, 0, 1) is equivalent to a.T.

Argument Description
input The input tensor
dim0 First dimension to swap
dim1 Second dimension to swap 
torch.transpose(a, 0, 1)  # shape (3, 4) → (4, 3)
print("Matrix a:")
print(a)
print(f"Shape of a: {a.shape}")
print("-"*79)

print("Transpose of a:")
print(a.T)
print(f"Shape of a.T: {a.T.shape}")
print("-"*79)

print("Transpose of a using torch.transpose:")
at = torch.transpose(a, 0, 1)
print(at)
print(f"Shape of at: {at.shape}")
Matrix a:
tensor([[-1.0349, -0.7509,  0.2990, -0.3119],
        [-0.2913,  0.0147,  0.7607,  0.7079],
        [-0.0297, -0.7357, -0.7373,  0.1570]])
Shape of a: torch.Size([3, 4])
-------------------------------------------------------------------------------
Transpose of a:
tensor([[-1.0349, -0.2913, -0.0297],
        [-0.7509,  0.0147, -0.7357],
        [ 0.2990,  0.7607, -0.7373],
        [-0.3119,  0.7079,  0.1570]])
Shape of a.T: torch.Size([4, 3])
-------------------------------------------------------------------------------
Transpose of a using torch.transpose:
tensor([[-1.0349, -0.2913, -0.0297],
        [-0.7509,  0.0147, -0.7357],
        [ 0.2990,  0.7607, -0.7373],
        [-0.3119,  0.7079,  0.1570]])
Shape of at: torch.Size([4, 3])

Note: For tensors with more than 2 dimensions, .T reverses all axes (e.g. (2, 3, 4) → (4, 3, 2)). PyTorch recommends using .mT (transposes only the last two dimensions) or .permute() for higher-dimensional tensors to avoid ambiguity.

# 3D tensor with shape (2, 3, 4)
x = torch.arange(24).reshape(2, 3, 4)
print("Original tensor x:")
print(x)
print(f"Original shape: {x.shape}")
print("-"*79)

# .T reverses ALL axes: (2, 3, 4) → (4, 3, 2)
print("Tranpose x with .T:")
print(x.T)
print(f"x.T shape:      {x.T.shape}")
print("-"*79)

# .mT transposes only the LAST TWO dimensions: (2, 3, 4) → (2, 4, 3)
print("Transpose x with .mT:")
print(x.mT)
print(f"x.mT shape:     {x.mT.shape}")
print("-"*79)

# .permute() gives full control over axis order: (2, 3, 4) → (4, 2, 3)
print("Transpose x with .permute(2, 0, 1):")
print(x.permute(2, 0, 1))
print(f"x.permute(2, 0, 1) shape: {x.permute(2, 0, 1).shape}")
Original tensor x:
tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])
Original shape: torch.Size([2, 3, 4])
-------------------------------------------------------------------------------
Tranpose x with .T:
tensor([[[ 0, 12],
         [ 4, 16],
         [ 8, 20]],

        [[ 1, 13],
         [ 5, 17],
         [ 9, 21]],

        [[ 2, 14],
         [ 6, 18],
         [10, 22]],

        [[ 3, 15],
         [ 7, 19],
         [11, 23]]])
x.T shape:      torch.Size([4, 3, 2])
-------------------------------------------------------------------------------
Transpose x with .mT:
tensor([[[ 0,  4,  8],
         [ 1,  5,  9],
         [ 2,  6, 10],
         [ 3,  7, 11]],

        [[12, 16, 20],
         [13, 17, 21],
         [14, 18, 22],
         [15, 19, 23]]])
x.mT shape:     torch.Size([2, 4, 3])
-------------------------------------------------------------------------------
Transpose x with .permute(2, 0, 1):
tensor([[[ 0,  4,  8],
         [12, 16, 20]],

        [[ 1,  5,  9],
         [13, 17, 21]],

        [[ 2,  6, 10],
         [14, 18, 22]],

        [[ 3,  7, 11],
         [15, 19, 23]]])
x.permute(2, 0, 1) shape: torch.Size([4, 2, 3])
/tmp/ipykernel_1001412/4065663704.py:10: UserWarning: The use of `x.T` on tensors of dimension other than 2 to reverse their shape is deprecated and it will throw an error in a future release. Consider `x.mT` to transpose batches of matrices or `x.permute(*torch.arange(x.ndim - 1, -1, -1))` to reverse the dimensions of a tensor. (Triggered internally at /__w/pytorch/pytorch/aten/src/ATen/native/TensorShape.cpp:4314.)
  print(x.T)

Stacking and Splitting Tensors

Stacking Tensors Vertically with torch.vstack()

  • torch.vstack(tensors) concatenates a sequence of tensors along the first dimension (axis 0, i.e. row-wise). All tensors must have the same shape in every dimension except dimension 0.
  • torch.vstack is equivalent to torch.cat(tensors, dim=0) for tensors with 2 or more dimensions.
Argument Description
tensors A tuple or list of tensors to stack. All must have the same shape except along dimension 0

Example:

a = torch.randn(3, 4)       # shape (3, 4)
b = torch.randn(3, 4)       # shape (3, 4)
c = torch.vstack((a, b))    # shape (6, 4)

The rows of b are appended below the rows of a, producing a tensor with 3 + 3 = 6 rows and the same 4 columns.

a = torch.randn(3,4)
print(a)
print(f"Shape of a: {a.shape}")
print("-"*79)

b = torch.randn(3,4)
print(b)
print(f"Shape of b: {b.shape}")
print("-"*79)
    
c = torch.vstack((a,b))
print("Stacked tensor c:")
print(c)
print(f"Shape of c: {c.shape}")
tensor([[ 0.2357, -0.6817, -1.4878, -1.5862],
        [-0.1810, -0.9178, -0.2779,  0.7317],
        [-0.5184, -0.8004, -0.4294,  0.7096]])
Shape of a: torch.Size([3, 4])
-------------------------------------------------------------------------------
tensor([[ 0.7678, -0.2684, -0.8338,  1.7119],
        [-1.0349, -1.0360, -0.8327,  0.9656],
        [ 0.2147,  1.5543, -0.0938,  0.2978]])
Shape of b: torch.Size([3, 4])
-------------------------------------------------------------------------------
Stacked tensor c:
tensor([[ 0.2357, -0.6817, -1.4878, -1.5862],
        [-0.1810, -0.9178, -0.2779,  0.7317],
        [-0.5184, -0.8004, -0.4294,  0.7096],
        [ 0.7678, -0.2684, -0.8338,  1.7119],
        [-1.0349, -1.0360, -0.8327,  0.9656],
        [ 0.2147,  1.5543, -0.0938,  0.2978]])
Shape of c: torch.Size([6, 4])

Stacking Tensors Horizontally with torch.hstack()

  • torch.hstack(tensors) concatenates a sequence of tensors along the second dimension (axis 1, i.e. column-wise). All tensors must have the same shape in every dimension except dimension 1.
  • torch.hstack is equivalent to torch.cat(tensors, dim=1) for tensors with 2 or more dimensions.
Argument Description
tensors A tuple or list of tensors to stack. All must have the same shape except along dimension 1

Example:

a = torch.randn(3, 4)    # shape (3, 4)
b = torch.randn(3, 4)    # shape (3, 4)
c = torch.hstack((a, b)) # shape (3, 8)

The columns of b are appended to the right of a, producing a tensor with the same 3 rows and 4 + 4 = 8 columns.

a = torch.randn(3,4)
print(a)
print(f"Shape of a: {a.shape}")
print("-"*79)

b = torch.randn(3,4)
print(b)
print(f"Shape of b: {b.shape}")
print("-"*79)

c = torch.hstack((a,b))
print("Stacked tensor c:")
print(c)
print(f"Shape of c: {c.shape}")
tensor([[-0.1912,  1.3198,  0.4368,  1.1204],
        [-0.1695,  0.1443, -1.2485,  0.8310],
        [ 0.8130, -0.6964,  1.3155,  0.3772]])
Shape of a: torch.Size([3, 4])
-------------------------------------------------------------------------------
tensor([[-1.4024, -1.6047,  0.4482,  0.1134],
        [-1.3010, -0.2858,  0.2347,  1.5640],
        [-0.4488,  0.9678,  1.0112, -0.9459]])
Shape of b: torch.Size([3, 4])
-------------------------------------------------------------------------------
Stacked tensor c:
tensor([[-0.1912,  1.3198,  0.4368,  1.1204, -1.4024, -1.6047,  0.4482,  0.1134],
        [-0.1695,  0.1443, -1.2485,  0.8310, -1.3010, -0.2858,  0.2347,  1.5640],
        [ 0.8130, -0.6964,  1.3155,  0.3772, -0.4488,  0.9678,  1.0112, -0.9459]])
Shape of c: torch.Size([3, 8])

Splitting Tensors Vertically with torch.vsplit()

  • torch.vsplit(input, sections) splits a tensor into multiple sub-tensors along the first dimension (axis 0, i.e. row-wise). It is the inverse of torch.vstack().
  • When sections is an integer, the size of dimension 0 must be evenly divisible by it. For example, splitting 6 rows into 2 sections means 3 rows per split (6 \div 2 = 3).
Argument Description
input The tensor to split
sections Either an int (number of equal splits) or a list of ints (sizes of each split along dim 0)

Example:

a = torch.randn(3, 4)        # shape (3, 4)
b = torch.randn(3, 4)        # shape (3, 4)
c = torch.vstack((a, b))     # shape (6, 4)
torch.vsplit(c, 2)           # → tuple of 2 tensors, each (3, 4)

The stacked tensor c has 6 rows. torch.vsplit(c, 2) divides it into 2 equal parts along dimension 0, giving two tensors of shape (3, 4), effectively recovering the original a and b.

a = torch.randn(3,4)
print(a)
print(f"Shape of a: {a.shape}")
print("-"*79)

b = torch.randn(3,4)
print(b)
print("-"*79)
print(f"Shape of b: {b.shape}")

c = torch.vstack((a,b))
print("Stacked tensor c:")
print(c)
print(f"Shape of c: {c.shape}")
print("-"*79)

print("Splitting c into 2 tensors along the vertical axis:")
split_tensor = torch.vsplit(c, 2)
print(split_tensor)
tensor([[ 0.5625, -0.1344,  2.5721,  0.9146],
        [-0.3667,  0.2697,  0.0829, -1.6199],
        [-0.6811, -0.5431,  0.9466, -0.4978]])
Shape of a: torch.Size([3, 4])
-------------------------------------------------------------------------------
tensor([[-0.5466,  0.3632,  0.7699, -0.1665],
        [-0.1693, -0.8723, -0.4792, -0.6744],
        [-1.5130, -0.0344, -0.7421,  1.6522]])
-------------------------------------------------------------------------------
Shape of b: torch.Size([3, 4])
Stacked tensor c:
tensor([[ 0.5625, -0.1344,  2.5721,  0.9146],
        [-0.3667,  0.2697,  0.0829, -1.6199],
        [-0.6811, -0.5431,  0.9466, -0.4978],
        [-0.5466,  0.3632,  0.7699, -0.1665],
        [-0.1693, -0.8723, -0.4792, -0.6744],
        [-1.5130, -0.0344, -0.7421,  1.6522]])
Shape of c: torch.Size([6, 4])
-------------------------------------------------------------------------------
Splitting c into 2 tensors along the vertical axis:
(tensor([[ 0.5625, -0.1344,  2.5721,  0.9146],
        [-0.3667,  0.2697,  0.0829, -1.6199],
        [-0.6811, -0.5431,  0.9466, -0.4978]]), tensor([[-0.5466,  0.3632,  0.7699, -0.1665],
        [-0.1693, -0.8723, -0.4792, -0.6744],
        [-1.5130, -0.0344, -0.7421,  1.6522]]))

Splitting Tensors by Chunk Size with torch.split()

  • torch.split(tensor, split_size_or_sections, dim=0) splits a tensor into chunks along a given dimension. Unlike vsplit which takes the number of equal splits, torch.split takes the size of each chunk.
  • If the tensor cannot be evenly divided, the last chunk will be smaller.
  • torch.split works along any dimension via the dim argument, making it more general than vsplit.
Argument Description
tensor The tensor to split
split_size_or_sections An int (size of each chunk) or a list of ints (exact size of each chunk)
dim The dimension along which to split (default: 0)

torch.split vs torch.vsplit

torch.split(c, 4, dim=0) torch.vsplit(c, 2)
Second argument means Chunk size (4 rows each) Number of equal splits
Result for (6, 4) 2 tensors of shape (4, 4) and (2, 4) 2 tensors of shape (3, 4)

Example:

torch.split(c, 4, dim=0)

The tensor c has shape (6, 4). torch.split(c, 4, dim=0) splits it along dimension 0 into chunks of 4 rows each, producing 2 tensors of shape (4, 4) and (2, 4). Since 6 is not evenly divisible by 4, the last chunk contains only the remaining 2 rows.

print("Original tensor c:")
c = torch.randn(6,4)
print(c)
print(f"Shape of c: {c.shape}")
print("-"*79)

print("Splitting c into 4 chunks per tensor using torch.split:")
split_tensor = torch.split(c, 4, dim=0)
print(f"Number of tensors after split: {len(split_tensor)}")
for i, t in enumerate(split_tensor):
    print(f"Tensor {i}:\n{t}\nShape: {t.shape}\n")
Original tensor c:
tensor([[-1.5944,  0.5426, -0.8320, -0.2036],
        [ 0.3461,  0.8811, -0.8798,  0.4232],
        [-0.9452,  0.0220, -0.1102,  0.3376],
        [ 0.1969,  1.0937, -0.5713, -2.0258],
        [-0.3106,  0.0613,  0.8255, -0.4412],
        [-1.7448,  1.0958, -0.8617, -0.4800]])
Shape of c: torch.Size([6, 4])
-------------------------------------------------------------------------------
Splitting c into 4 chunks per tensor using torch.split:
Number of tensors after split: 2
Tensor 0:
tensor([[-1.5944,  0.5426, -0.8320, -0.2036],
        [ 0.3461,  0.8811, -0.8798,  0.4232],
        [-0.9452,  0.0220, -0.1102,  0.3376],
        [ 0.1969,  1.0937, -0.5713, -2.0258]])
Shape: torch.Size([4, 4])

Tensor 1:
tensor([[-0.3106,  0.0613,  0.8255, -0.4412],
        [-1.7448,  1.0958, -0.8617, -0.4800]])
Shape: torch.Size([2, 4])