NumPy Matrix Operations

Real-World Scenario

A machine learning engineer implements a neural network forward pass from scratch to debug a training issue. The forward pass is pure matrix multiplication: input matrix × weight matrix + bias vector, repeated for each layer. Understanding NumPy matrix operations means understanding the mathematical backbone of every deep learning framework.

Matrix Multiplication

import numpy as np

# Matrix multiplication — the @ operator (preferred) or np.matmul
A = np.array([[1, 2], [3, 4]])   # 2×2
B = np.array([[5, 6], [7, 8]])   # 2×2

C = A @ B
print(C)
# [[19 22]
#  [43 50]]
# C[i,j] = sum(A[i,:] * B[:,j]) — dot product of row i with column j

# Rectangular matrices: (m×k) @ (k×n) → (m×n)
X = np.random.default_rng(0).standard_normal((100, 20))  # 100 samples, 20 features
W = np.random.default_rng(1).standard_normal((20, 10))   # 20 inputs, 10 outputs
b = np.zeros(10)

output = X @ W + b   # (100, 20) @ (20, 10) + (10,) → (100, 10)
print(output.shape)  # (100, 10)

# Dot product of two vectors (1-D arrays)
u = np.array([1., 2., 3.])
v = np.array([4., 5., 6.])
print(np.dot(u, v))  # 32.0 = 1*4 + 2*5 + 3*6

Transpose

import numpy as np

A = np.array([[1, 2, 3], [4, 5, 6]])  # shape (2, 3)

# .T attribute — transpose rows and columns
A_T = A.T
print(A_T.shape)  # (3, 2)
print(A_T)
# [[1 4]
#  [2 5]
#  [3 6]]

# Transpose is a view — modifying it modifies the original
A_T[0, 0] = 99
print(A[0, 0])  # 99

# np.transpose with explicit axes — useful for N-D tensors
# Swap axes 0 and 2 of a 3-D array (batch, height, width) → (width, height, batch)
tensor = np.ones((8, 28, 28))        # 8 images, 28×28 pixels
reordered = np.transpose(tensor, (2, 1, 0))  # (28, 28, 8)
print(reordered.shape)  # (28, 28, 8)

Reshape and Flatten

import numpy as np

arr = np.arange(24)  # [0, 1, 2, ..., 23]

# reshape — change shape without changing data
matrix = arr.reshape(4, 6)   # 4 rows, 6 columns
print(matrix.shape)  # (4, 6)

# Use -1 to let NumPy infer one dimension
cube = arr.reshape(2, 3, -1)  # 2 × 3 × ? — NumPy computes ? = 4
print(cube.shape)  # (2, 3, 4)

# Common ML pattern: flatten all spatial dims, keep batch dim
images = np.ones((32, 28, 28))     # 32 images of 28×28 pixels
flat = images.reshape(32, -1)      # (32, 784) — flatten each image
print(flat.shape)  # (32, 784)

# flatten() always returns a 1-D copy
flat_copy = matrix.flatten()
print(flat_copy.shape)  # (24,)

# ravel() returns a 1-D view when possible (more efficient than flatten)
flat_view = matrix.ravel()
print(flat_view.shape)  # (24,)

Stacking and Splitting

import numpy as np

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

# Stack along a new axis (creates a new dimension)
stacked = np.stack([a, b])          # (2, 3) — stacks as rows
col_stack = np.stack([a, b], axis=1)  # (3, 2) — stacks as columns

# Concatenate along an existing axis
hstack = np.hstack([a, b])          # [1 2 3 4 5 6] — horizontal (axis=1 for 2D)
print(hstack.shape)  # (6,)

A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])

vstack = np.vstack([A, B])          # stack rows (axis=0)
print(vstack.shape)  # (4, 2)

hstack2d = np.hstack([A, B])        # stack columns (axis=1)
print(hstack2d.shape)  # (2, 4)

# Split — inverse of stack/concatenate
arr = np.arange(12).reshape(3, 4)

# Split into 3 equal parts along axis=1 (columns)
parts = np.hsplit(arr, 3)           # 3 arrays of shape (3, 1)... actually (3,4/3)
# Uneven splits use array of split points
left, right = np.hsplit(arr, [2])   # split at column index 2
print(left.shape, right.shape)      # (3, 2) (3, 2)

Real-World: Neural Network Forward Pass

import numpy as np

def relu(x: np.ndarray) -> np.ndarray:
    """ReLU activation: max(0, x)."""
    return np.maximum(0, x)

def softmax(x: np.ndarray) -> np.ndarray:
    """Numerically stable softmax."""
    # Subtract max per sample to prevent overflow
    exp_x = np.exp(x - x.max(axis=1, keepdims=True))
    return exp_x / exp_x.sum(axis=1, keepdims=True)

rng = np.random.default_rng(42)

# Network: 784 inputs → 256 hidden → 128 hidden → 10 outputs
layer_sizes = [784, 256, 128, 10]

# Initialize weights with He initialization, biases at zero
weights = [
    rng.standard_normal((layer_sizes[i], layer_sizes[i+1])) * np.sqrt(2 / layer_sizes[i])
    for i in range(len(layer_sizes) - 1)
]
biases = [np.zeros(layer_sizes[i+1]) for i in range(len(layer_sizes) - 1)]

# Forward pass for a batch of 32 flattened MNIST images
batch_size = 32
X = rng.standard_normal((batch_size, 784))   # input batch

activations = X
for W, b in zip(weights[:-1], biases[:-1]):
    activations = relu(activations @ W + b)  # linear + activation

# Final layer — softmax for probabilities
logits = activations @ weights[-1] + biases[-1]
probs = softmax(logits)

print(f"Input shape:  {X.shape}")       # (32, 784)
print(f"Output shape: {probs.shape}")   # (32, 10)
print(f"Row sums:     {probs.sum(axis=1)[:3].round(6)}")  # [1. 1. 1.]

Element-wise vs Matrix Multiply

import numpy as np

A = np.array([[1., 2.], [3., 4.]])
B = np.array([[2., 0.], [1., 3.]])

# Element-wise multiplication — the Hadamard product
elementwise = A * B
print(elementwise)
# [[2. 0.]
#  [3. 12.]]

# Matrix multiplication
matmul = A @ B
print(matmul)
# [[ 4.  6.]
#  [10. 12.]]