Matrix Multiplication in Python: From Nested Loops to NumPy and PyTorch

The Foundation of Modern Computing

Matrix multiplication is the mathematical heartbeat of computer graphics, robotics physics simulations, cryptographic transformations, and deep learning attention heads. Understanding how matrices interact at the algorithmic level enables developers to write optimized numerical algorithms and diagnose computational bottlenecks in data science pipelines.

To multiply two matrices A and B, the number of columns in A must strictly match the number of rows in B. If matrix A has dimensions (M × K) and matrix B has dimensions (K × N), the resulting matrix C will have dimensions (M × N).

Method 1: Pure Python with Triple Nested Loops

Writing matrix multiplication in raw Python without libraries illustrates the fundamental dot-product arithmetic. Each cell C[i][j] is the sum of products of row i from A and column j from B:

def manual_matrix_multiply(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    rows_A = len(A)
    cols_A = len(A[0])
    rows_B = len(B)
    cols_B = len(B[0])

    # Validation: inner dimensions must match
    if cols_A != rows_B:
        raise ValueError(f"Dimension mismatch: cannot multiply ({rows_A}x{cols_A}) with ({rows_B}x{cols_B})")

    # Initialize result matrix C with dimensions (rows_A x cols_B)
    C = [[0.0 for _ in range(cols_B)] for _ in range(rows_A)]

    # Triple nested loop: O(N^3) time complexity
    for i in range(rows_A):
        for j in range(cols_B):
            cell_sum = 0.0
            for k in range(cols_A):
                cell_sum += A[i][k] * B[k][j]
            C[i][j] = cell_sum

    return C

# Example execution
A = [[1, 2, 3],
     [4, 5, 6]]

B = [[7, 8],
     [9, 1],
     [2, 3]]

result = manual_matrix_multiply(A, B)
# Output: [[31.0, 19.0], [85.0, 55.0]]
print("Manual Output:", result)

While intuitive, this triple-loop implementation suffers from significant Python interpreter overhead and poor CPU cache locality, making it unusable for matrices larger than a few hundred rows.

Method 2: Modern NumPy with the @ Infix Operator

NumPy bypasses Python interpreter loops by compiling matrix arithmetic down to optimized C and Fortran linear algebra libraries (BLAS/LAPACK) utilizing SIMD vector registers.

Modern Python (3.5+) provides the dedicated matrix multiplication infix operator: @.

import numpy as np

# Define matrices as typed NumPy arrays
matrix_A = np.array([
    [1, 2, 3],
    [4, 5, 6]
], dtype=np.float64)

matrix_B = np.array([
    [7, 8],
    [9, 1],
    [2, 3]
], dtype=np.float64)

# Modern, clean syntax using the @ operator
result_infix = matrix_A @ matrix_B

# Equivalent explicit NumPy function call
result_matmul = np.matmul(matrix_A, matrix_B)

print("NumPy Infix Output:\n", result_infix)
# Output shape: (2, 2)

Method 3: GPU Acceleration with PyTorch

When training neural networks or processing batch embeddings, CPUs cannot compete with the thousands of parallel tensor cores on modern GPUs. PyTorch allows you to move matrices directly to GPU VRAM and perform parallel hardware multiplication:

import torch

# Detect available hardware accelerator (CUDA or Apple Silicon MPS)
device = torch.device("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"))

# Allocate random 2000x2000 tensors directly on acceleration device
tensor_A = torch.randn(2000, 2000, device=device, dtype=torch.float32)
tensor_B = torch.randn(2000, 2000, device=device, dtype=torch.float32)

# High-speed GPU tensor multiplication
tensor_C = tensor_A @ tensor_B

print(f"Computed on {device}: resulting shape {tensor_C.shape}")

Performance Benchmarks: 1000x1000 Matrix Multiplication

To quantify the performance disparity between these approaches, consider multiplying two square matrices of size 1000 × 1000 (which requires one billion floating-point operations):

• Pure Python Triple Loops: Approximately 85 to 120 seconds. Memory cache misses and dynamic type-checking on every multiplication create extreme drag.
• NumPy with OpenBLAS: Approximately 0.018 seconds (over 5,000 times faster than pure Python) thanks to multi-threading and cache-blocked tiling.
• PyTorch with CUDA on NVIDIA RTX: Approximately 0.0012 seconds (sub-2-millisecond latency) using parallel FP32 tensor cores.

Common Pitfalls and How to Avoid Them

Keep these three practical rules in mind during linear algebra development:

• Element-Wise Multiplication vs Matrix Product: In Python NumPy, A * B performs element-by-element Hadamard multiplication, requiring identical dimensions. Always use A @ B for algebraic dot-product matrix multiplication.
• Memory Contiguity: NumPy operations are fastest when arrays are C-contiguous in memory. If you transpose an array before multiplication, use np.ascontiguousarray(A.T) if profiling indicates memory stalls.
• Floating-Point Precision: Be conscious of precision. Multiplying large matrices in float32 (single precision) provides double the throughput of float64 (double precision) on modern hardware, which is why deep learning models standardise on float32 or bfloat16.

Summary

Matrix multiplication is an indispensable computer science primitive. While vanilla nested loops help solidify mathematical understanding, production applications in data science and web engineering should always use the native @ operator via NumPy or PyTorch to benefit from vectorized hardware acceleration.