Einsum Fundamentals - Einstein Summation (Einsum)

Einsum Fundamentals

What you'll build

A function einsum_basics(arr) that computes five single-matrix operations (total sum, row sums, column sums, transpose and diagonal) using only np.einsum.

Background

np.einsum(subscripts, *operands) describes an operation by labelling each axis with a letter. The rules:

• letters after -> are kept, in that order;
• letters that appear in the input but not after -> are summed over;
• a letter repeated within one operand ('ii') walks the diagonal.

Starting with one operand makes the rules easy to see before they're used for matrix products and attention.

Subscripts for a 2D array arr of shape (m, n)

Subscripts	Result	Shape	Equivalent
'ij->'	sum of all entries	scalar	arr.sum()
'ij->i'	sum of each row (j is summed)	(m,)	arr.sum(axis=1)
'ij->j'	sum of each column (i is summed)	(n,)	arr.sum(axis=0)
'ij->ji'	transpose	(n, m)	arr.T
'ii->i'	diagonal (square only)	(m,)	np.diag(arr)

Inputs

• arr: 2D numeric array of shape (m, n); it may or may not be square.

Expected Return Format

Return a dictionary with these keys:

• 'sum_all': sum of all elements (scalar)
• 'row_sum': sum of each row, shape (m,)
• 'col_sum': sum of each column, shape (n,)
• 'transpose': the transposed array, shape (n, m)
• 'diagonal': the diagonal, shape (m,), if arr is square; otherwise None

Edge cases and conventions

• 'ii->i' needs both axes to have the same size; on a non-square array it raises an error, so check arr.shape[0] == arr.shape[1] first and return None otherwise.
• Integer input gives integer results.

Constraints

Compute every value with np.einsum (that's the point of the exercise).

Follow-up

• What does 'ii->' compute? And np.einsum('ji', arr), with no arrow?

Input: einsum_basics(np.array([[1, 2], [3, 4]]))
Output: {'sum_all': 10, 'row_sum': array([3, 7]), 'col_sum': array([4, 6]), 'transpose': array([[1, 3], [2, 4]]), 'diagonal': array([1, 4])}

Input: einsum_basics(np.array([[1, 2, 3], [4, 5, 6]]))
Output: {'sum_all': 21, 'row_sum': array([ 6, 15]), 'col_sum': array([5, 7, 9]), 'transpose': array([[1, 4], [2, 5], [3, 6]]), 'diagonal': None}

import numpy as np

def einsum_basics(arr: np.ndarray) -> dict:
    """
    Single-matrix operations with np.einsum.

    Args:
        arr: 2D array of shape (m, n)

    Returns:
        dict with keys:
            'sum_all':   sum of all elements (scalar)
            'row_sum':   row sums, shape (m,)
            'col_sum':   column sums, shape (n,)
            'transpose': arr transposed, shape (n, m)
            'diagonal':  diagonal, shape (m,), if arr is square, else None
    """
    # Your code here
    pass