Naive matrix multiplication in Python
Compute each entry of the product AB as the dot product of a row of A and a column of B, with three nested loops and n m p multiplications.
- File
naive_matrix_multiplication.py, 30 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/naive-matrix-multiplication/. The authors are listed in prov/F9X6P7.jsonl.
Matrix = list[list[int]]
def naive_matrix_multiplication(a: Matrix, b: Matrix, p: int | None = None, mod: int | None = None) -> Matrix:
"""Return the product a * b by the triple loop.
a is n x m and b is m x p, given as lists of rows. When m = 0, b is empty,
so the caller passes the number of columns p. With mod set, the arithmetic
is in Z/mod Z and each entry is in [0, mod). Cost: n*m*p multiplications
and n*m*p additions; extra space O(1) beyond the output.
"""
n = len(a)
m = len(b)
if p is None:
p = len(b[0]) if m > 0 else 0
if any(len(row) != m for row in a) or any(len(row) != p for row in b):
raise ValueError("the shapes do not match")
c = [[0] * p for _ in range(n)] # @line L1
for i in range(n): # @line L2
for j in range(p): # @line L3
s = 0 # @line L4
for k in range(m): # @line L5
s = s + a[i][k] * b[k][j] # @line L6
if mod is not None:
s %= mod
c[i][j] = s # @line L7
return c # @line L8
test_naive_matrix_multiplication.py 82 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/naive-matrix-multiplication/. The authors are listed in prov/F9X6P7.jsonl.
import random
import pytest
from naive_matrix_multiplication import naive_matrix_multiplication as mul
def oracle(a, b, p, mod=None):
"""An independent product: each column of b is formed first, then zipped with each row of a."""
cols = [[row[j] for row in b] for j in range(p)]
out = [[sum(x * y for x, y in zip(r, col)) for col in cols] for r in a]
if mod is not None:
out = [[x % mod for x in r] for r in out]
return out
def rand(rng, r, c, lo=-9, hi=9):
return [[rng.randint(lo, hi) for _ in range(c)] for _ in range(r)]
def test_example() -> None:
a = [[1, 2, 3], [4, 5, 6]]
b = [[7, 8], [9, 10], [11, 12]]
assert mul(a, b) == [[58, 64], [139, 154]]
def test_edges() -> None:
assert mul([], []) == []
assert mul([[5]], [[7]]) == [[35]]
assert mul([[], []], [], p=3) == [[0, 0, 0], [0, 0, 0]]
assert mul([[1, 2]], [[3], [4]]) == [[11]]
assert mul([[3], [4]], [[1, 2]]) == [[3, 6], [4, 8]]
with pytest.raises(ValueError):
mul([[1, 2]], [[1, 2]])
def test_identity_and_zero() -> None:
rng = random.Random(11)
for n in range(0, 7):
a = rand(rng, n, n)
eye = [[int(i == j) for j in range(n)] for i in range(n)]
zero = [[0] * n for _ in range(n)]
assert mul(a, eye, p=n) == a
assert mul(eye, a, p=n) == a
assert mul(a, zero, p=n) == zero
def test_random_over_z() -> None:
rng = random.Random(2024)
for _ in range(600):
n, m, p = rng.randint(0, 6), rng.randint(0, 6), rng.randint(0, 6)
a, b = rand(rng, n, m, -1000, 1000), rand(rng, m, p, -1000, 1000)
assert mul(a, b, p=p) == oracle(a, b, p)
def test_random_mod_p() -> None:
rng = random.Random(7)
for q in (2, 3, 5, 7, 101):
for _ in range(100):
n, m, p = rng.randint(1, 6), rng.randint(0, 6), rng.randint(1, 6)
a, b = rand(rng, n, m, 0, q - 1), rand(rng, m, p, 0, q - 1)
assert mul(a, b, p=p, mod=q) == oracle(a, b, p, q)
def test_laws() -> None:
rng = random.Random(99)
for _ in range(100):
n, m, p, r = (rng.randint(1, 5) for _ in range(4))
a, b, c = rand(rng, n, m), rand(rng, m, p), rand(rng, p, r)
assert mul(mul(a, b), c) == mul(a, mul(b, c))
bt = [list(x) for x in zip(*b)]
at = [list(x) for x in zip(*a)]
assert [list(x) for x in zip(*mul(a, b))] == mul(bt, at)
def test_large_integers() -> None:
big = 2**63 - 1
a = [[big, big], [big, big]]
assert mul(a, a) == [[2 * big * big] * 2] * 2
What the next level needs
- the tests did not run in this build