Matrix Multiplication with Einsum

easyFREEEinsum OperationsPyTorch Advanced

Problem Statement

Use torch.einsum to perform matrix multiplication.

Background

Einstein summation (einsum) is a compact notation for tensor operations. The key idea is that repeated indices are summed over, while indices that appear in the output are preserved.

Your Task

The starter code creates matrices A (2x3) and B (3x2). Use torch.einsum to compute the matrix product of A and B.

The verification against torch.matmul is pre-filled.

Output Format

Returns a dictionary with "result", "shape", and "matches_matmul".

Example

Input

None

Output

{'result': [[58.0, 64.0], [139.0, 154.0]], 'shape': [2, 2], 'matches_matmul': True}

Reasoning

Inputs: two fixed float matrices, $A = \begin{bmatrix} 1 & 2 & 3 \\ 4 & 5 & 6 \end{bmatrix}$ ( $2 \times 3$ ) and $B = \begin{bmatrix} 7 & 8 \\ 9 & 10 \\ 11 & 12 \end{bmatrix}$ ( $3 \times 2$ ).
Math: label $A$ 's axes $i, j$ and $B$ 's axes $j, k$ . The shared index $j$ (length 3) is summed away; $i$ and $k$ survive: $C_{ik} = \sum_j A_{ij} B_{jk}$ .
Row 1 of $A$ against each column of $B$ : $1 \cdot 7 + 2 \cdot 9 + 3 \cdot 11 = 58$ and $1 \cdot 8 + 2 \cdot 10 + 3 \cdot 12 = 64$ .
Row 2: $4 \cdot 7 + 5 \cdot 9 + 6 \cdot 11 = 139$ and $4 \cdot 8 + 5 \cdot 10 + 6 \cdot 12 = 154$ .
Output: result is [[58.0, 64.0], [139.0, 154.0]] (floats, like the inputs), shape is [2, 2] because only $i$ and $k$ remain, and matches_matmul is True since this index rule is ordinary matrix multiplication.
Notice: the summed index must have equal length (3) in both matrices, which is why $2 \times 3$ times $3 \times 2$ works.

Constraints

Use torch.einsum with 'ij,jk->ik'
Verify against torch.matmul
Return as nested list

Learn this in

PyTorch FundamentalsFoundations

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 primer

Background Knowledge

Introduction to Einsum

Einstein summation (einsum) is a mathematical notation used to describe tensor operations in a compact and expressive way. It is based on the concept of implicit summation over repeated indices. In the context of matrix multiplication, einsum provides a concise way to describe the operation without explicitly writing out the loops. The notation 'ij,jk->ik' represents the matrix multiplication operation, where i and k are the output indices, and j is the shared index that is summed over.

Matrix Multiplication

Matrix multiplication is a fundamental operation in linear algebra, where two matrices are multiplied to produce another matrix. Given two matrices A and B, the matrix product C = A @ B is defined as the sum of the products of the elements of each row of A and each column of B. In this problem, we are given two matrices A (2x3) and B (3x2), and we need to compute the matrix product C using torch.einsum.

PyTorch and Einsum

PyTorch provides an implementation of einsum through the torch.einsum function, which allows us to perform tensor operations using the einsum notation. This function takes the einsum equation as a string, followed by the input tensors. In this problem, we will use torch.einsum to perform matrix multiplication and compare the result with the torch.matmul function.

Algorithm/Approach

The general approach to solving this problem involves:

Creating the input matrices A and B
Using torch.einsum to perform matrix multiplication
Verifying the result using torch.matmul
Returning the result and its shape, along with a flag indicating whether the einsum result matches the torch.matmul result

Step-by-Step Strategy

To implement the solution, follow these steps:

Import the necessary PyTorch module and create the input matrices A and B.
Use torch.einsum to perform matrix multiplication, specifying the einsum equation and the input tensors.
Compute the matrix product using torch.matmul for verification.
Compare the results of torch.einsum and torch.matmul to determine if they match.
Return a dictionary containing the einsum result, its shape, and a flag indicating whether the results match.

Common Pitfalls

When implementing the solution, watch out for:

Incorrect einsum equation: Make sure to specify the correct einsum equation for matrix multiplication.
Input tensor shapes: Verify that the input tensors have the correct shapes.
Data types: Ensure that the input tensors have the correct data type (float).

Time & Space Complexity

The time complexity of the torch.einsum operation is O(n^3), where n is the number of elements in the input tensors. The space complexity is O(n^2), where n is the number of elements in the output tensor. The torch.matmul operation has the same time and space complexity. However, the actual performance may vary depending on the specific implementation and hardware.

solution.py

Python 3.13

Test results