Tensor Contraction with Einsum

mediumPROEinsum OperationsPyTorch Advanced

Problem Statement

Use einsum for advanced tensor contractions: Frobenius inner product, row-wise dot products, and column sums.

Background

Indices that appear on the left but not the right are summed over. By controlling which indices appear in the output, you can express a variety of reductions and contractions.

Your Task

The starter code creates matrices A and B. Use einsum to compute three operations: the Frobenius inner product (element-wise multiply and sum all), row-wise dot products (sum along columns only), and column sums of A.

Output Format

Returns a dictionary with "frobenius", "row_dots", and "col_sums".

Example

Input

None

Output

{'frobenius': 70.0, 'row_dots': [17.0, 53.0], 'col_sums': [4.0, 6.0]}

Reasoning

Inputs: the function builds two fixed $2 \times 2$ float matrices, $A = \begin{pmatrix}1 & 2\\ 3 & 4\end{pmatrix}$ and $B = \begin{pmatrix}5 & 6\\ 7 & 8\end{pmatrix}$ . Index $i$ picks the row and $j$ the column, so $A_{ij}$ and $B_{ij}$ sit in the same cell.
Frobenius inner product: multiply matching cells and sum over both indices, $\sum_{i,j} A_{ij}B_{ij} = 5 + 12 + 21 + 32 = 70$ . No index survives, so frobenius is a single scalar, 70.0.
Row-wise dot products: keep $i$ and sum over $j$ only, $r_i = \sum_j A_{ij}B_{ij}$ . Row 0 gives $1 \cdot 5 + 2 \cdot 6 = 17$ , row 1 gives $3 \cdot 7 + 4 \cdot 8 = 53$ , so row_dots is [17.0, 53.0], one value per row.
Column sums of A: only $A$ is involved; keep $j$ and sum over $i$ , $c_j = \sum_i A_{ij}$ , giving $1 + 3 = 4$ and $2 + 4 = 6$ , so col_sums is [4.0, 6.0].
Notice that the row dots add up to the Frobenius value ( $17 + 53 = 70$ ): summing out $i$ afterwards finishes the full contraction. The indices you keep decide the output's shape.

Constraints

'ij,ij->' for Frobenius inner product
'ij,ij->i' for row-wise dots
'ij->j' for column sums

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

Concept primerPREVIEW

Background Knowledge

The problem involves using einsum for advanced tensor contractions. Einsum is a powerful operation in linear algebra that allows for concise and expressive representation of various tensor operations, including contractions. It is based on the Einstein summation convention, which implies summation over repeated indices. In the context of tensors, einsum can be used to perform a wide range of operations, from simple element-wise multiplication to complex multi-axis contractions.

To understand einsum, it's essential to grasp the basics of tensor operations and the Einstein summation convention. Tensors are multi-dimensional arrays, and operations on them can be described using indices. The Einstein summation convention states that when an index appears twice in a term, it implies summation over that index. For example, the expression 'ij,ij->' represents the element-wise product sum of two tensors, where i and j are the indices of the tensors. The -> symbol indicates the output shape, which in this case is a scalar.

In PyTorch, einsum is implemented as a function torch.einsum(), which takes a string describing the operation and the input tensors. The string describes the indices of the input tensors and the output shape. For instance, 'ij,ij->' would perform an element-wise product sum of two 2D tensors, resulting in a scalar. Understanding the syntax and semantics of einsum is crucial for solving this problem.

Algorithm/Approach

The general approach to solving this problem involves using einsum to perform the required tensor contractions. The algorithm can be broken down into the following steps:

Define the input tensors A and B
Use einsum to compute the element-wise product sum, row-wise dot products, and column sums
Store the results in a dictionary and return it

The key concept here is to understand how to use einsum to describe the desired tensor operations. This involves choosing the correct indices and output shape for each operation.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results