Einsum vs Traditional Operations - Einstein Summation (Einsum)

Einsum vs Traditional Operations

What you'll build

A function einsum_equivalence(A, B) that computes five operations twice, once with np.einsum and once with the usual NumPy call, and reports whether every pair matches.

Background

In code reviews and interviews you'll need to translate fluently in both directions: read 'ik,kj->ij' as A @ B, and write A.sum(axis=1) as 'ij->i'. Checking the two against each other is also how you debug an einsum you're unsure of.

The five operations

Key	Einsum	NumPy equivalent
'sum'	np.einsum('ij->', A)	np.sum(A)
'row_sum'	np.einsum('ij->i', A)	np.sum(A, axis=1)
'matmul'	np.einsum('ik,kj->ij', A, B)	A @ B
'transpose'	np.einsum('ij->ji', A)	A.T
'hadamard_sum'	np.einsum('ij,ij->', A, A)	np.sum(A * A)

Note that 'hadamard_sum' multiplies A with itself (the squared Frobenius norm of A); B is only used for 'matmul'.

Inputs

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

Expected Return Format

Return a dictionary with two keys:

• 'results': a dict with exactly the five keys 'sum', 'row_sum', 'matmul', 'transpose', 'hadamard_sum'. Each value is a dict {'einsum': <einsum result>, 'numpy': <NumPy result>}.
• 'all_match': a Python bool, True if every einsum result is close to its NumPy counterpart (np.allclose).

Edge cases and conventions

• Compare with np.allclose, not ==: floating-point sums can differ in the last bits depending on summation order.
• all(...) over the comparisons returns a Python bool; np.all would return np.bool_.

Constraints

NumPy only.

Follow-up

• Which of these einsums is a view and which creates a new array?
• For 'ik,kj->ij' on large matrices, which is faster in practice, einsum or @, and why?

Input: r = einsum_equivalence(np.array([[1., 2., 3.], [4., 5., 6.]]), np.array([[1., 2.], [3., 4.], [5., 6.]]))
r["all_match"]
Output: True

Input: r["results"]["matmul"]["einsum"]
Output: array([[22., 28.],
       [49., 64.]])

Input: r["results"]["hadamard_sum"]
Output: {'einsum': 91.0, 'numpy': 91.0}

import numpy as np

def einsum_equivalence(A: np.ndarray, B: np.ndarray) -> dict:
    """
    Compare einsum with traditional NumPy operations.

    Args:
        A: First matrix (m, k)
        B: Second matrix (k, n)

    Returns:
        dict with keys:
            'results':   {'sum', 'row_sum', 'matmul', 'transpose', 'hadamard_sum'},
                         each mapping to {'einsum': ..., 'numpy': ...}
            'all_match': Python bool, True if every pair is np.allclose
    """
    # Your code here
    pass