Tensors in Action: From Element-wise Math to Matrix Multiplication in PyTorch

This notebook provides a hands-on guide to tensor operations in PyTorch, from basic element-wise math to matrix multiplication.

• What are the common element-wise tensor operations?
• How do statistical and aggregation operations work across dimensions?
• How can we detect NaN values with torch.isnan()?
• What is the vector dot product and how does torch.dot() work?
• How does matrix multiplication work with torch.matmul() and the @ operator?
• How does singular value decomposition work with torch.linalg.svd()?
• How can we visualise a 3D saddle surface using torch.meshgrid and matplotlib?

Common Element-wise Tensor Operations

Element-wise operations apply a function independently to each element of a tensor, producing an output tensor of the same shape. These operations are vectorised in PyTorch, meaning they run efficiently on the entire tensor without explicit loops. The table below summarises some of the most commonly used element-wise functions.

Function	What it does
torch.exp(a)	Computes e^x for each element
torch.log(a)	Computes the natural logarithm (\ln x) for each element
torch.sin(a)	Computes the sine of each element (in radians)
torch.arctan(a)	Computes the inverse tangent (arctan) of each element
torch.abs(a)	Returns the absolute value of each element
torch.square(a)	Squares each element (x^2)
torch.sqrt(a)	Computes the square root of each element (\sqrt{x})
torch.ceil(a)	Rounds each element up to the nearest integer
torch.round(a)	Rounds each element to the nearest integer
torch.clip(a, min, max)	Clamps each element to the range [min, max]

import torch
import matplotlib.pyplot as plt

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

print("exp =", torch.exp(a))
print("-"*79)

print("log =", torch.log(a))
print("-"*79)

print("sin =", torch.sin(a))
print("-"*79)

print("arctan =", torch.arctan(a))
print("-"*79)

Original tensor a:
tensor([[ 1.2598, -0.0041,  1.3176],
        [-0.5203, -0.5858,  0.0835]])
-------------------------------------------------------------------------------
exp = tensor([[3.5248, 0.9959, 3.7344],
        [0.5944, 0.5567, 1.0870]])
-------------------------------------------------------------------------------
log = tensor([[ 0.2310,     nan,  0.2758],
        [    nan,     nan, -2.4833]])
-------------------------------------------------------------------------------
sin = tensor([[ 0.9520, -0.0041,  0.9681],
        [-0.4971, -0.5529,  0.0834]])
-------------------------------------------------------------------------------
arctan = tensor([[ 0.8999, -0.0041,  0.9216],
        [-0.4797, -0.5299,  0.0833]])
-------------------------------------------------------------------------------

print("abs =", torch.abs(a))
print("-"*79)

print("square =", torch.square(a))
print("-"*79)

print("sqrt =", torch.sqrt(a))
print("-"*79)

print("ceil =", torch.ceil(a))
print("-"*79)

print("round =", torch.round(a))
print("-"*79)

print("clip =", torch.clip(a, 0.1, 0.9))

abs = tensor([[1.2598, 0.0041, 1.3176],
        [0.5203, 0.5858, 0.0835]])
-------------------------------------------------------------------------------
square = tensor([[1.5871e+00, 1.6864e-05, 1.7360e+00],
        [2.7068e-01, 3.4314e-01, 6.9662e-03]])
-------------------------------------------------------------------------------
sqrt = tensor([[1.1224,    nan, 1.1479],
        [   nan,    nan, 0.2889]])
-------------------------------------------------------------------------------
ceil = tensor([[2., -0., 2.],
        [-0., -0., 1.]])
-------------------------------------------------------------------------------
round = tensor([[ 1., -0.,  1.],
        [-1., -1.,  0.]])
-------------------------------------------------------------------------------
clip = tensor([[0.9000, 0.1000, 0.9000],
        [0.1000, 0.1000, 0.1000]])

Statistical and Aggregation Operations

Statistical and aggregation operations reduce a tensor along one or more dimensions, summarising its values into means, sums, extrema, and other descriptive statistics. All functions below accept an optional dim argument to operate along a specific dimension; without dim, they reduce over the entire tensor. The table below lists the most commonly used aggregation functions in PyTorch.

Function	What it does
torch.mean(a, dim)	Computes the arithmetic mean along a dimension
torch.std(a, dim)	Computes the standard deviation along a dimension
torch.var(a, dim)	Computes the variance along a dimension
torch.median(a, dim)	Returns the median value (and indices) along a dimension
torch.sum(a, dim)	Sums all elements along a dimension
torch.prod(a, dim)	Computes the product of all elements along a dimension
torch.cumsum(a, dim)	Computes the cumulative sum along a dimension
torch.cumprod(a, dim)	Computes the cumulative product along a dimension
torch.min(a, dim)	Returns the minimum value (and indices) along a dimension
torch.max(a, dim)	Returns the maximum value (and indices) along a dimension
torch.argmin(a, dim)	Returns the index of the minimum value along a dimension
torch.argmax(a, dim)	Returns the index of the maximum value along a dimension
torch.norm(a, dim)	Computes the L^p norm along a dimension (default p=2)
torch.quantile(a, q, dim)	Computes the q-th quantile along a dimension

For a 2D tensor with shape (m, n), setting dim=0 reduces along the rows (the first axis), collapsing m rows into a single row. The result has shape (n,), which means one value per column, computed from all m entries in that column. For example, torch.mean(a, dim=0) on a (3, 4) tensor returns 4 values: the mean of each column across the 3 rows.

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

print("mean (dim=0):")
print(torch.mean(a, dim=0))
print("-"*79)

print("std (dim=0):")
print(torch.std(a, dim=0))
print("-"*79)

print("var (dim=0):")
print(torch.var(a, dim=0))
print("-"*79)

print("median (dim=0):")
print(torch.median(a, dim=0))
print("-"*79)

print("sum (dim=0):")
print(torch.sum(a, dim=0))
print("-"*79)

print("prod (dim=0):")
print(torch.prod(a, dim=0))
print("-"*79)

Tensor a:
tensor([[-0.2427,  0.2384, -0.1037,  0.8094],
        [ 0.1877, -1.7584,  1.7207, -2.3809],
        [ 0.6550, -0.2082,  1.6822,  0.6771]])
-------------------------------------------------------------------------------
mean (dim=0):
tensor([ 0.2000, -0.5761,  1.0997, -0.2981])
-------------------------------------------------------------------------------
std (dim=0):
tensor([0.4490, 1.0480, 1.0424, 1.8050])
-------------------------------------------------------------------------------
var (dim=0):
tensor([0.2016, 1.0983, 1.0865, 3.2579])
-------------------------------------------------------------------------------
median (dim=0):
torch.return_types.median(
values=tensor([ 0.1877, -0.2082,  1.6822,  0.6771]),
indices=tensor([1, 2, 2, 2]))
-------------------------------------------------------------------------------
sum (dim=0):
tensor([ 0.6000, -1.7282,  3.2992, -0.8944])
-------------------------------------------------------------------------------
prod (dim=0):
tensor([-0.0298,  0.0873, -0.3001, -1.3049])
-------------------------------------------------------------------------------

print("cumsum (dim=0):")
print(torch.cumsum(a, dim=0))
print("-"*79)

print("cumprod (dim=0):")
print(torch.cumprod(a, dim=0))
print("-"*79)

print("min (dim=0):")
print(torch.min(a, dim=0))
print("-"*79)

print("max (dim=0):")
print(torch.max(a, dim=0))
print("-"*79)

print("argmin (dim=0):")
print(torch.argmin(a, dim=0))
print("-"*79)

print("argmax (dim=0):")
print(torch.argmax(a, dim=0))
print("-"*79)

print("norm (dim=0):")
print(torch.norm(a, dim=0))
print("-"*79)

print("quantile (q=0.5, dim=0):")
print(torch.quantile(a, 0.5, dim=0))

cumsum (dim=0):
tensor([[-0.2427,  0.2384, -0.1037,  0.8094],
        [-0.0550, -1.5200,  1.6170, -1.5715],
        [ 0.6000, -1.7282,  3.2992, -0.8944]])
-------------------------------------------------------------------------------
cumprod (dim=0):
tensor([[-0.2427,  0.2384, -0.1037,  0.8094],
        [-0.0456, -0.4192, -0.1784, -1.9271],
        [-0.0298,  0.0873, -0.3001, -1.3049]])
-------------------------------------------------------------------------------
min (dim=0):
torch.return_types.min(
values=tensor([-0.2427, -1.7584, -0.1037, -2.3809]),
indices=tensor([0, 1, 0, 1]))
-------------------------------------------------------------------------------
max (dim=0):
torch.return_types.max(
values=tensor([0.6550, 0.2384, 1.7207, 0.8094]),
indices=tensor([2, 0, 1, 0]))
-------------------------------------------------------------------------------
argmin (dim=0):
tensor([0, 1, 0, 1])
-------------------------------------------------------------------------------
argmax (dim=0):
tensor([2, 0, 1, 0])
-------------------------------------------------------------------------------
norm (dim=0):
tensor([0.7233, 1.7866, 2.4086, 2.6043])
-------------------------------------------------------------------------------
quantile (q=0.5, dim=0):
tensor([ 0.1877, -0.2082,  1.6822,  0.6771])

Detecting NaN Values with torch.isnan()

torch.isnan(input) returns a boolean tensor of the same shape, where each element is True if the corresponding value is NaN (Not a Number), and False otherwise.

Why NaN appears here

torch.sqrt(a) computes the square root of each element. Since a was created with torch.randn(), some elements may be negative. The square root of a negative real number is undefined, so PyTorch returns NaN for those entries.

b = torch.sqrt(a)
print(b)
print(torch.isnan(b))

tensor([[   nan, 0.4883,    nan, 0.8997],
        [0.4333,    nan, 1.3117,    nan],
        [0.8093,    nan, 1.2970, 0.8229]])
tensor([[ True, False,  True, False],
        [False,  True, False,  True],
        [False,  True, False, False]])

Vector Dot Product with torch.dot()

• The dot product, named after the dot symbol (\cdot) used to denote it, is an operation on vectors that takes two vectors and returns a single number (a scalar), unlike addition or scalar multiplication which return vectors.
• The dot product multiplies corresponding elements and sums the results, producing a single scalar:

a \cdot b = a_1 b_1 + a_2 b_2 + a_3 b_3 = \sum_{i=1}^{n} a_i b_i

• It is also called the scalar product for this reason, and is a specific case of a more general operation known as the inner product.
• torch.dot(a, b) computes the dot product of two 1D tensors. The @ operator also works for 1D inputs and produces the same result.
• The result is a 0D scalar tensor (not a 1D tensor).
• The dot product is closely related to matrix multiplication: for 1D vectors, a @ b is equivalent to torch.dot(a, b), whereas for 2D matrices @ performs torch.matmul.

Example:

a = torch.randn(3)   # shape (3,)
b = torch.randn(3)   # shape (3,)
torch.dot(a, b)      # scalar
a @ b                # scalar

An Illustration of Vector Dot Product

a = torch.randn(3)
print("Tensor a:")
print(a)
print("-"*79)

b = torch.randn(3)
print("Tensor b:")
print(b)
print("-"*79)

print("Dot product using torch.dot:")
print(torch.dot(a, b))
print("-"*79)

print("Dot product using @ operator:")
print(a @ b)

Tensor a:
tensor([-2.1615,  1.2231,  2.0254])
-------------------------------------------------------------------------------
Tensor b:
tensor([1.3500, 0.6283, 0.7026])
-------------------------------------------------------------------------------
Dot product using torch.dot:
tensor(-0.7265)
-------------------------------------------------------------------------------
Dot product using @ operator:
tensor(-0.7265)

Matrix Multiplication with torch.matmul() and the @ Operator

• Matrix multiplication requires the inner dimensions to match: (m, k) × (k, n) → (m, n). Here a is (2, 4) and b is (4, 3), so the inner dimension is 4 and the result is (2, 3). Each element c_{ij} is the dot product of row i of a and column j of b:

c_{ij} = \sum_{k=1}^{4} a_{ik} \cdot b_{kj}

For example:

• c_{11} = a_{11}b_{11} + a_{12}b_{21} + a_{13}b_{31} + a_{14}b_{41} (row 1 of a \cdot column 1 of b)
• c_{22} = a_{21}b_{12} + a_{22}b_{22} + a_{23}b_{32} + a_{24}b_{42} (row 2 of a \cdot column 2 of b)

An Illustration of Matrix Multiplication

• torch.matmul(a, b) performs matrix multiplication between two tensors. The @ operator is a shorthand that does the same thing.
• This is not element-wise multiplication. For that, use a * b or torch.mul(a, b).

torch.matmul vs @ vs other options

Method	Description
torch.matmul(a, b)	General-purpose matrix multiplication; supports broadcasting and batched inputs
a @ b	Operator shorthand for torch.matmul
torch.mm(a, b)	Strictly 2D matrix multiplication only (no broadcasting)

Example:

a = torch.randn(2, 4)   # shape (2, 4)
b = torch.randn(4, 3)   # shape (4, 3)
torch.matmul(a, b)      # shape (2, 3)
a @ b                   # shape (2, 3)

a = torch.randn(2, 4)
print("Tensor a (2x4):")
print(a)
print("-"*79)

b = torch.randn(4, 3)
print("Tensor b (4x3):")
print(b)
print("-"*79)

print("Matrix multiplication using torch.matmul (2x3):")
print(torch.matmul(a, b))
print("-"*79)

print("Matrix multiplication using @ operator (2x3):")
print(a @ b)

Tensor a (2x4):
tensor([[ 1.2194,  1.3448,  0.2941,  1.6100],
        [ 0.1784, -0.3189, -1.0860,  0.6163]])
-------------------------------------------------------------------------------
Tensor b (4x3):
tensor([[-2.2816,  1.3947, -0.4564],
        [-0.3457,  0.7439, -1.4234],
        [ 1.5466,  1.4533,  0.9812],
        [ 0.6984,  0.5244,  0.6321]])
-------------------------------------------------------------------------------
Matrix multiplication using torch.matmul (2x3):
tensor([[-1.6678,  3.9728, -1.1646],
        [-1.5460, -1.2436, -0.3036]])
-------------------------------------------------------------------------------
Matrix multiplication using @ operator (2x3):
tensor([[-1.6678,  3.9728, -1.1646],
        [-1.5460, -1.2436, -0.3036]])

Singular Value Decomposition (SVD) with torch.linalg.svd()

What is SVD?

• SVD breaks any matrix into three simpler matrices: a rotation (U), a scaling (diagonal S), and another rotation (V^H), so that A = U \cdot \text{diag}(S) \cdot V^H.
• The singular values in S tell you how much each corresponding direction is stretched, ranked from most to least important. The first value captures the most variance.
• This makes SVD a powerful tool for tasks like dimensionality reduction and low-rank approximation, since you can drop the smallest singular values to get a compact approximation of the original matrix.
• U and Vh are orthogonal matrices, meaning U^T U = I and V^H V = I.
• torch.linalg.svd(a) computes the Singular Value Decomposition of a matrix A, factorizing it into three matrices:

A = U \cdot \text{diag}(S) \cdot V^H

Example:

a = torch.randn(3, 3)   # shape (3, 3)
torch.linalg.svd(a)     # returns (U, S, Vh)

The output is a named tuple with three components:

Component	Shape	Meaning
U	(m, m) → (3, 3)	Left singular vectors: orthogonal matrix whose columns are eigenvectors of AA^T
S	(min(m, n),) → (3,)	Singular values: non-negative values in descending order
Vh	(n, n) → (3, 3)	Right singular vectors (conjugate-transposed): orthogonal matrix whose rows are eigenvectors of A^T A

Verifying: Reconstruct A from U, S, and V^H

We can manually verify the decomposition by reconstructing the original matrix from its SVD components. If the decomposition is correct, U \cdot \text{diag}(S) \cdot V^H should equal the original matrix A (up to floating-point precision).

a = torch.randn(3, 3)
print("Tensor a:")
print(a)
print("-"*79)

# Compute SVD
U, S, Vh = torch.linalg.svd(a)
print("U:")
print(U)
print("-"*79)

print("S:")
print(S)
print("-"*79)

print("Vh:")
print(Vh)
print("-"*79)

# Reconstruct A from U, S, Vh
a_reconstructed = U @ torch.diag(S) @ Vh
print("Reconstructed A = U @ diag(S) @ Vh:")
print(a_reconstructed)
print("-"*79)

# Verify: difference should be near zero
print("Difference (A - reconstructed A):")
print(a - a_reconstructed)
print("-"*79)
print("All close?", torch.allclose(a, a_reconstructed))

Tensor a:
tensor([[ 0.6534,  1.1631, -1.6241],
        [ 1.0289, -2.0669,  1.3471],
        [ 0.8866,  0.8535, -0.8044]])
-------------------------------------------------------------------------------
U:
tensor([[ 0.5794,  0.5063, -0.6387],
        [-0.7382,  0.6581, -0.1480],
        [ 0.3454,  0.5573,  0.7550]])
-------------------------------------------------------------------------------
S:
tensor([3.3357, 1.5783, 0.3258])
-------------------------------------------------------------------------------
Vh:
tensor([[-0.0224,  0.7478, -0.6635],
        [ 0.9517, -0.1874, -0.2433],
        [ 0.3063,  0.6369,  0.7075]])
-------------------------------------------------------------------------------
Reconstructed A = U @ diag(S) @ Vh:
tensor([[ 0.6534,  1.1631, -1.6241],
        [ 1.0289, -2.0669,  1.3471],
        [ 0.8866,  0.8535, -0.8044]])
-------------------------------------------------------------------------------
Difference (A - reconstructed A):
tensor([[-1.1921e-07,  2.3842e-07, -2.3842e-07],
        [-1.1921e-07,  0.0000e+00,  0.0000e+00],
        [-1.7881e-07,  1.1921e-07, -5.9605e-08]])
-------------------------------------------------------------------------------
All close? True

Manual SVD Computation without torch.linalg.svd

The SVD of a real matrix A can be derived from the eigendecomposition of A^T A:

1. Compute A^T A: a symmetric positive semi-definite matrix of shape (n, n)
2. Eigendecompose A^T A = V \Lambda V^T, where \Lambda = \text{diag}(\lambda_1, \lambda_2, \ldots)
3. The singular values are \sigma_i = \sqrt{\lambda_i}, sorted in descending order
4. The columns of V (reordered by descending \sigma_i) give V^H = V^T
5. Compute each column of U as: u_i = \frac{1}{\sigma_i} A v_i

# Use the same tensor a from the previous cell
print("Original tensor a:")
print(a)
print("-"*79)

# Step 1: Compute A^T A
AtA = a.T @ a
print("Step 1: Compute A^T A")
print(AtA)
print("-"*79)

# Step 2: Eigendecompose A^T A - eigenvalues and eigenvectors
eigenvalues, V = torch.linalg.eigh(AtA)
print("Step 2: Eigenvalues of A^T A")
print(eigenvalues)
print("Eigenvectors V (columns):")
print(V)
print("-"*79)

# Step 3: Sort eigenvalues in descending order
print("Step 3: Sort eigenvalues in descending order")
print("Original eigenvalues:", eigenvalues)
sort_idx = torch.argsort(eigenvalues, descending=True)
print("Sorted indices:", sort_idx)
eigenvalues = eigenvalues[sort_idx]
print("Sorted eigenvalues:", eigenvalues)
V = V[:, sort_idx]
print("Sorted eigenvectors V (columns):")
print(V)
print("-"*79)

# Step 4: Singular values = sqrt(eigenvalues)
S_manual = torch.sqrt(eigenvalues)
print("Step 4: Singular values S (sorted descending):")
print(S_manual)
print("-"*79)

# Step 5: Vh = V^T
Vh_manual = V.T
print("Step 5: Vh = V^T:")
print(Vh_manual)
print("-"*79)

Original tensor a:
tensor([[ 0.6534,  1.1631, -1.6241],
        [ 1.0289, -2.0669,  1.3471],
        [ 0.8866,  0.8535, -0.8044]])
-------------------------------------------------------------------------------
Step 1: Compute A^T A
tensor([[ 2.2716, -0.6098, -0.3885],
        [-0.6098,  6.3533, -5.3598],
        [-0.3885, -5.3598,  5.0993]])
-------------------------------------------------------------------------------
Step 2: Eigenvalues of A^T A
tensor([ 0.1062,  2.4910, 11.1271])
Eigenvectors V (columns):
tensor([[ 0.3063,  0.9517,  0.0224],
        [ 0.6369, -0.1874, -0.7478],
        [ 0.7075, -0.2433,  0.6635]])
-------------------------------------------------------------------------------
Step 3: Sort eigenvalues in descending order
Original eigenvalues: tensor([ 0.1062,  2.4910, 11.1271])
Sorted indices: tensor([2, 1, 0])
Sorted eigenvalues: tensor([11.1271,  2.4910,  0.1062])
Sorted eigenvectors V (columns):
tensor([[ 0.0224,  0.9517,  0.3063],
        [-0.7478, -0.1874,  0.6369],
        [ 0.6635, -0.2433,  0.7075]])
-------------------------------------------------------------------------------
Step 4: Singular values S (sorted descending):
tensor([3.3357, 1.5783, 0.3258])
-------------------------------------------------------------------------------
Step 5: Vh = V^T:
tensor([[ 0.0224, -0.7478,  0.6635],
        [ 0.9517, -0.1874, -0.2433],
        [ 0.3063,  0.6369,  0.7075]])
-------------------------------------------------------------------------------

# Step 6: Compute U column by column: u_i = (1/σ_i) * A @ v_i
U_manual = torch.zeros(a.shape[0], a.shape[0])
for i in range(len(S_manual)):
    U_manual[:, i] = (a @ V[:, i]) / S_manual[i]
print("Step 6: U (computed as u_i = A v_i / σ_i):")
print(U_manual)
print("-"*79)

# Verify: reconstruct A from manual SVD
a_manual = U_manual @ torch.diag(S_manual) @ Vh_manual
print("Reconstructed A = U @ diag(S) @ Vh:")
print(a_manual)
print("-"*79)
print("All close to original?", torch.allclose(a, a_manual, atol=1e-5))

Step 6: U (computed as u_i = A v_i / σ_i):
tensor([[-0.5794,  0.5063, -0.6387],
        [ 0.7382,  0.6581, -0.1480],
        [-0.3454,  0.5573,  0.7550]])
-------------------------------------------------------------------------------
Reconstructed A = U @ diag(S) @ Vh:
tensor([[ 0.6534,  1.1631, -1.6241],
        [ 1.0289, -2.0669,  1.3471],
        [ 0.8866,  0.8535, -0.8044]])
-------------------------------------------------------------------------------
All close to original? True

Visualising a Hyperbolic Paraboloid (Saddle Surface)

This section demonstrates how to construct and visualise a 3D mathematical surface from several PyTorch tensor operations. The surface plotted is a hyperbolic paraboloid (z = x^2 - (y/2)^2), which curves upward in one direction and downward in another, forming a classic saddle shape. Cross-section lines and a marked saddle point at the origin highlight the opposing curvatures. It’s a practical example of how vectorised tensor operations can efficiently evaluate a function over an entire grid without any explicit Python loops.

How to Visualise a Surface Plot with PyTorch?

1. Create 1D grids: torch.linspace(-1, 1, 100) produces 100 evenly spaced values for x, and similarly for y over [-2, 2].
2. Expand to 2D grids: torch.meshgrid(x, y, indexing="xy") broadcasts the 1D vectors into two (100, 100) tensors xx and yy, where every (x, y) pair on the grid is represented.
3. Compute z element-wise: The surface equation z = x^2 - (y/2)^2 is evaluated for all 10 000 grid points at once using element-wise ** (power) and - (subtraction). This is a hyperbolic paraboloid: it curves upward along the x-axis and downward along the y-axis, meeting at a saddle point at the origin.
4. Plot the surface: ax.plot_surface renders the 3D surface with a diverging RdBu colour map (red for negative z, blue for positive z).
5. Overlay cross-sections: Two black lines through the saddle point show the opposing curvatures:
   • Solid (y = 0): z = x^2, which is an upward parabola.
   • Dashed (x = 0): z = -(y/2)^2, which is a downward parabola.
6. Mark the saddle point: A red dot at (0, 0, 0) highlights where the two curvatures meet.

Note: A saddle point is a point on a surface where the curvature goes up in one direction and down in another (like the centre of a horse saddle). At that point, the surface is neither a local minimum nor a local maximum; it’s a minimum along one axis, but a maximum along the other.

In this example, the origin (0, 0, 0) is the saddle point: along the x-axis the parabola z = x^2 curves upward, while along the y-axis the parabola z = -(y/2)^2 curves downward. Saddle points are important in optimisation (e.g. training neural networks) because gradient-based methods can slow down or get stuck near them, since the gradient is zero but the point is not a true minimum.

# create tensors
x = torch.linspace(-1, 1, 100)
y = torch.linspace(-2, 2, 100)

# create the surface
xx, yy = torch.meshgrid(
    x, y, indexing="xy")
print("Meshgrid xx:")
print(xx)
print("-"*79)
print("Meshgrid yy:")
print(yy)

# Define a hyperbolic paraboloid surface: z = x^2 - (y/2)^2
z = xx**2 - (yy/2)**2

# xy-indexing is matching numpy
fig = plt.figure(figsize=(8,8))
ax = fig.add_subplot(111, projection="3d")
ax.set_xlim([-2, 2])
ax.set_ylim([-2, 2])
ax.set_zlim([-2, 2])
ax.set_xlabel("x-axis", fontweight="bold", fontsize=12, labelpad=1)
ax.set_ylabel("y-axis", fontweight="bold", fontsize=12, labelpad=1)
ax.set_zlabel("z-axis", fontweight="bold", fontsize=12, labelpad=0.5)
ax.plot_surface(xx, yy, z, cmap="RdBu", alpha=0.8)

# Plot lines through the saddle point at (0, 0, 0)
# Along x-axis (y=0): z = x^2  (upward parabola)
ax.plot(
    x, torch.zeros_like(x), x**2,
    color="black", linewidth=3,
    label="y=0: z = x²")

# Along y-axis (x=0): z = -(y/2)^2  (downward parabola)
ax.plot(
    torch.zeros_like(y), y, -(y/2)**2,
    color="black", linewidth=3,
    linestyle="--", label="x=0: z = -(y/2)²")

# Mark the saddle point
ax.scatter(
    [0], [0], [0], color="red", s=100, zorder=5,
    label="Saddle point (0,0,0)")

ax.legend(fontsize=10)
# ax.view_init(45, 35)
plt.tight_layout()
plt.show()

Meshgrid xx:
tensor([[-1.0000, -0.9798, -0.9596,  ...,  0.9596,  0.9798,  1.0000],
        [-1.0000, -0.9798, -0.9596,  ...,  0.9596,  0.9798,  1.0000],
        [-1.0000, -0.9798, -0.9596,  ...,  0.9596,  0.9798,  1.0000],
        ...,
        [-1.0000, -0.9798, -0.9596,  ...,  0.9596,  0.9798,  1.0000],
        [-1.0000, -0.9798, -0.9596,  ...,  0.9596,  0.9798,  1.0000],
        [-1.0000, -0.9798, -0.9596,  ...,  0.9596,  0.9798,  1.0000]])
-------------------------------------------------------------------------------
Meshgrid yy:
tensor([[-2.0000, -2.0000, -2.0000,  ..., -2.0000, -2.0000, -2.0000],
        [-1.9596, -1.9596, -1.9596,  ..., -1.9596, -1.9596, -1.9596],
        [-1.9192, -1.9192, -1.9192,  ..., -1.9192, -1.9192, -1.9192],
        ...,
        [ 1.9192,  1.9192,  1.9192,  ...,  1.9192,  1.9192,  1.9192],
        [ 1.9596,  1.9596,  1.9596,  ...,  1.9596,  1.9596,  1.9596],
        [ 2.0000,  2.0000,  2.0000,  ...,  2.0000,  2.0000,  2.0000]])

Figure 1: An example of hyperbolic paraboloid surface plot with PyTorch.

References

1. PyTorch Documentation, 2024. Tensors. PyTorch.
2. PyTorch Documentation, 2024. torch.matmul. PyTorch.