Matrix Operations with Einsum - Einstein Summation (Einsum)

Matrix Operations with Einsum

What you'll build

A function einsum_matrix_ops(A, B, v) that computes six two-operand products (matrix multiply, matrix-vector, outer, Hadamard, Frobenius inner product and trace) with np.einsum.

Background

With two operands, a letter shared between them pairs up their axes. If that letter is missing from the output, the paired axis is contracted (multiplied and summed), which is how a matrix product works. If it's kept, the entries are just multiplied elementwise. Every linear layer, attention score and bilinear model is one of these patterns.

Formulas

Key	Subscripts	Formula	Shape
matmul	'ik,kj->ij'	C[i, j] = sum_k A[i, k] * B[k, j]	(m, n)
matvec	'ij,j->i'	y[i] = sum_j A[i, j] * v[j]	(m,)
outer_product	'i,j->ij'	O[i, j] = v[i] * v[j]	(k, k)
hadamard	'ij,ij->ij'	H[i, j] = A[i, j] * A[i, j]	(m, k)
frobenius	'ij,ij->'	sum_ij A[i, j] * A[i, j]	scalar
trace	'ii->'	sum_i A[i, i]	scalar

Inputs

• A: matrix of shape (m, k)
• B: matrix of shape (k, n)
• v: vector of shape (k,)

Expected Return Format

Return a dictionary with these keys:

• 'matmul': A @ B, shape (m, n)
• 'matvec': A @ v, shape (m,)
• 'outer_product': outer product of v with itself, shape (k, k)
• 'hadamard': elementwise A * A, shape (m, k)
• 'frobenius': Frobenius inner product of A with itself, (A * A).sum() (the squared Frobenius norm), a scalar
• 'trace': trace of A if A is square, otherwise None

Edge cases and conventions

• 'ii->' requires a square A; return None for the trace when m != k.
• Integer inputs give integer outputs.

Constraints

Compute every value with np.einsum.

Follow-up

• Write u^T A w (a bilinear form) as a single einsum with three operands.
• What is the cost of 'ik,kj->ij' in terms of m, k and n?

Input: einsum_matrix_ops(np.array([[1, 2], [3, 4]]), np.array([[5, 6], [7, 8]]), np.array([1, 2]))
Output: {'matmul': array([[19, 22], [43, 50]]), 'matvec': array([ 5, 11]), 'outer_product': array([[1, 2], [2, 4]]), 'hadamard': array([[ 1,  4], [ 9, 16]]), 'frobenius': 30, 'trace': 5}

Input: einsum_matrix_ops(np.array([[1, 0, 2]]), np.array([[1], [2], [3]]), np.array([1, 1, 1]))["trace"]
Output: None

import numpy as np

def einsum_matrix_ops(A: np.ndarray, B: np.ndarray, v: np.ndarray) -> dict:
    """
    Two-operand products with np.einsum.

    Args:
        A: Matrix of shape (m, k)
        B: Matrix of shape (k, n)
        v: Vector of shape (k,)

    Returns:
        dict with keys:
            'matmul':        A @ B, shape (m, n)
            'matvec':        A @ v, shape (m,)
            'outer_product': outer(v, v), shape (k, k)
            'hadamard':      A * A, shape (m, k)
            'frobenius':     (A * A).sum(), scalar
            'trace':         trace of A if square, else None
    """
    # Your code here
    pass