Matrix Multiplication - Python Basics for ML

Matrix Multiplication

What you'll build

A function matrix_multiply(A, B) that returns the matrix product of two 2D arrays.

Background

Matrix multiplication is the workhorse of deep learning: a linear layer computes x @ W, attention computes Q @ K.T, and GPUs are built to do it fast. Knowing the shape rule by heart lets you debug most "shapes don't match" errors in seconds.

Formula

For A of shape (m, n) and B of shape (n, p):

C[i, j] = sum over k = 0 .. n-1 of A[i, k] * B[k, j]        C has shape (m, p)

Entry (i, j) is the dot product of row i of A with column j of B. The inner dimensions must match and disappear.

Inputs

• A: 2D array of shape (m, n)
• B: 2D array of shape (n, p), with 1 <= m, n, p <= 100. Integer or float entries.

Expected Return Format

A 2D array of shape (m, p). Integer inputs give an integer result.

Edge cases and conventions

• A * B is the elementwise product, not matrix multiplication.
• A (1, n) @ (n, 1) product is still 2D: [[value]], shape (1, 1).

Constraints

Use a NumPy matrix-product routine (@, np.matmul or np.dot), not Python loops.

Follow-up

• Write the triple loop by hand. What is its time complexity? (O(m·n·p).)
• How do np.dot and np.matmul differ on 3D (batched) inputs?

Input: matrix_multiply(np.array([[1, 2], [3, 4]]), np.array([[5, 6], [7, 8]]))
Output: array([[19, 22],
       [43, 50]])

Input: matrix_multiply(np.array([[1, 2, 3]]), np.array([[1], [2], [3]]))
Output: array([[14]])

import numpy as np

def matrix_multiply(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """
    Multiply two matrices A and B.

    Args:
        A: First matrix of shape (m, n)
        B: Second matrix of shape (n, p)

    Returns:
        Matrix product A @ B of shape (m, p)
    """
    # Your code here
    pass