Trace and Diagonal with Einsum

easyPROEinsum OperationsPyTorch Advanced

Problem Statement

Use einsum to compute the trace, extract the diagonal, and compute an outer product.

Background

Repeated indices on the same tensor access diagonal elements. Summing repeated indices gives the trace, keeping them gives the diagonal. Two separate indices create an outer product.

Your Task

The starter code creates a matrix M and vectors a, b. Use einsum to compute three operations: the trace of M, the diagonal of M, and the outer product of a and b.

Output Format

Returns a dictionary with "trace", "diagonal", "outer_product", and "trace_equals_diag_sum".

Example

Input

None

Output

{'trace': 15.0, 'diagonal': [1.0, 5.0, 9.0], 'outer_product': [[4.0, 5.0], [8.0, 10.0], [12.0, 15.0]], 'trace_equals_diag_sum': True}

Reasoning

Inputs: a $3 \times 3$ matrix $M$ with rows $(1, 2, 3)$ , $(4, 5, 6)$ , $(7, 8, 9)$ , and two vectors $a = (1, 2, 3)$ of length 3 and $b = (4, 5)$ of length 2. All values are floats, hence the .0 endings.
Trace: the diagonal entries $M_{ii}$ share one index for row and column. Summing over $i$ gives a scalar: $\operatorname{tr}(M) = \sum_i M_{ii} = 1 + 5 + 9 = 15$ , giving 15.0.
Diagonal: the same entries, but with $i$ kept rather than summed away, so the result is a length-3 vector: [1.0, 5.0, 9.0].
Outer product: two different indices give every pairing, $(ab^\top)_{ij} = a_i b_j$ . With 3 values of $i$ and 2 of $j$ the result is $3 \times 2$ : row 1 is $1 \cdot (4, 5) = (4, 5)$ , row 2 is $(8, 10)$ , row 3 is $(12, 15)$ .
Check: $1 + 5 + 9 = 15$ matches the trace, so trace_equals_diag_sum is True. Notice that trace and diagonal differ only in whether the repeated index survives.

Constraints

'ii->' for trace
'ii->i' for diagonal
'i,j->ij' for outer product

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy