Batch Matrix Multiplication with Einsum

mediumPROEinsum OperationsPyTorch Advanced

Problem Statement

Use einsum for batched matrix multiplication.

Background

Batched operations perform independent computations for each element along the batch dimension. Einsum handles this naturally by including a batch index that is preserved (not summed over).

Your Task

The starter code creates batched tensors A (2x3x4) and B (2x4x5). Use torch.einsum to compute batched matrix multiplication.

The verification against torch.bmm is pre-filled.

Output Format

Returns a dictionary with "output_shape", "matches_bmm", "batch_size", and "first_element".

Example

Input

None

Output

{'output_shape': [2, 3, 5], 'matches_bmm': True, 'batch_size': 2, 'first_element': 3.8424}

Reasoning

Inputs: a fixed random seed produces two stacks of matrices: $A$ is $2 \times 3 \times 4$ (two $3 \times 4$ matrices) and $B$ is $2 \times 4 \times 5$ (two $4 \times 5$ matrices). The leading axis $b$ is the batch.
Math: each batch slice is multiplied on its own, $C_{bik} = \sum_{j} A_{bij} B_{bjk}$ . The shared index $j$ (length 4) is summed away, while $b$ , $i$ and $k$ survive, so $C$ is $2 \times 3 \times 5$ , giving output_shape [2, 3, 5] and batch_size 2.
Worked entry: $C_{000}$ pairs row 0 of the first $A$ matrix, $[1.9269, 1.4873, 0.9007, -2.1055]$ , with column 0 of the first $B$ matrix, $[0.0349, 0.6872, 0.0499, -1.2862]$ : $0.06727 + 1.02200 + 0.04495 + 2.70815 \approx 3.84238$ , rounded to 4 decimal places as 3.8424.
Check: a dedicated batched matrix product computes the same numbers, so matches_bmm is True.
Notice that the batch index is never summed: slice 0 of $A$ only ever meets slice 0 of $B$ .

Constraints

Use 'bij,bjk->bik' notation
Verify against torch.bmm
Batch dimension preserved

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 primerPREVIEW

Background Knowledge

The problem involves using PyTorch and its torch.einsum function to perform batched matrix multiplication. Einsum is a powerful notation for expressing multi-dimensional tensor operations. It allows for a concise and expressive way to describe complex tensor operations, including matrix multiplication. In the context of this problem, we're dealing with batched matrix multiplication, where we have multiple matrices (in this case, 2) that we want to multiply together in a batched fashion.

To understand this problem, it's essential to have a grasp of tensor operations and matrix multiplication. In PyTorch, tensors are multi-dimensional arrays that can be used to represent a wide range of data, including matrices. Matrix multiplication is a fundamental operation in linear algebra, where two matrices are multiplied together to produce another matrix. The torch.einsum function provides a flexible way to perform matrix multiplication and other tensor operations by specifying the operation using a string notation. For example, the string 'bij,bjk->bik' specifies a batched matrix multiplication operation, where b represents the batch dimension, i represents the rows of the first matrix, j represents the columns of the first matrix (and rows of the second matrix), and k represents the columns of the second matrix.

The torch.bmm function is another way to perform batched matrix multiplication in PyTorch. It takes two tensors as input and returns the result of the batched matrix multiplication. In this problem, we're asked to verify that the result of the torch.einsum operation matches the result of the torch.bmm operation.

Algorithm/Approach

The general approach to solving this problem involves using the torch.einsum function to perform batched matrix multiplication and then verifying the result against the torch.bmm function. The key concept here is to understand how to specify the batched matrix multiplication operation using the torch.einsum notation. This involves identifying the batch dimension and the dimensions of the matrices being multiplied.

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