np.einsum() for efficient computation in NumPy - Step-by-Step Execution

Concept Flow - np.einsum() for efficient computation

Input Arrays

↓

Define Subscripts

↓

np.einsum() Parses Subscripts

↓

Perform Summation/Product

↓

Return Result Array

np.einsum takes input arrays and a subscript string to specify element-wise operations and summations, then computes the result efficiently.

Execution Sample

NumPy

import numpy as np
A = np.array([[1,2],[3,4]])
B = np.array([[5,6],[7,8]])
result = np.einsum('ij,ij->ij', A, B)

Computes element-wise multiplication of two 2x2 arrays using np.einsum.

Execution Table

Step	Subscript Parsing	Operation	Intermediate Result	Output
1	'ij,ij->ij'	Parse indices: i,j for both arrays, output i,j	Arrays ready for element-wise multiply	None yet
2	'ij,ij->ij'	Element-wise multiplication	[[5,12],[21,32]]	None yet
3	No summation since output is 'ij'	Keep shape (2,2)	[[5,12],[21,32]]	[[5,12],[21,32]]

💡 Completed element-wise multiplication with output shape (2,2) as specified by 'ij' in subscript.

Variable Tracker

Variable	Start	After Step 2	Final
A	[[1,2],[3,4]]	[[1,2],[3,4]]	[[1,2],[3,4]]
B	[[5,6],[7,8]]	[[5,6],[7,8]]	[[5,6],[7,8]]
result	None	[Intermediate: element-wise multiply]	[[5,12],[21,32]]

Key Moments - 2 Insights

Why does the output keep the same shape as inputs?

Because the subscript 'ij,ij->ij' tells np.einsum to multiply elements with matching indices and keep the indices 'ij' in output, so no summation happens (see execution_table step 3).

What if we change the output subscript to 'i'?

np.einsum will sum over the 'j' index, reducing the output to 1D array with shape matching 'i'. This is a summation step different from element-wise multiply (not shown in current table).

Visual Quiz - 3 Questions

Look at the execution table, what is the intermediate result after step 2?

A [[1,2],[3,4]]

B [[6,8],[10,12]]

C [[5,12],[21,32]]

D None

💡 Hint

Check the 'Intermediate Result' column at step 2 in execution_table.

At which step does np.einsum parse the subscripts?

A Step 1

B Step 2

C Step 3

D No parsing step

💡 Hint

Look at the 'Subscript Parsing' column in execution_table.

If the output subscript changes to 'i', what happens to the output shape?

A Output shape becomes (2,2)

B Output shape becomes (2,)

C Output shape becomes (1,2)

D Output shape becomes scalar

💡 Hint

Refer to key_moments explanation about summation over 'j' index.

Concept Snapshot

np.einsum(subscripts, *arrays) computes sums/products using index notation.
Subscripts specify input and output indices.
Matching indices are multiplied element-wise.
Indices not in output are summed over.
Efficient for complex array operations without loops.

Full Transcript

np.einsum is a powerful numpy function that lets you specify how to multiply and sum elements of arrays using index labels. You write a string like 'ij,ij->ij' to tell it to multiply elements with matching i and j indices and keep the output shape the same. The function parses this string, multiplies elements accordingly, and sums over indices not in the output. This example showed element-wise multiplication of two 2x2 arrays, resulting in a 2x2 output. Changing the output indices can cause summation and change the output shape. This method is efficient and avoids explicit loops.

Practice

1. What is the main purpose of np.einsum() in numpy?

easy

A. To perform efficient and readable array operations using Einstein summation notation

B. To create random arrays with specific distributions

C. To sort arrays in ascending order

D. To reshape arrays into different dimensions

Solution

1. Step 1: Understand the function's goal

   np.einsum() uses Einstein summation notation to express complex array operations clearly and efficiently.

2. Step 2: Compare with other options

   Options B, C, and D describe unrelated numpy functions like random generation, sorting, and reshaping.

3. Final Answer:

   To perform efficient and readable array operations using Einstein summation notation -> Option A

4. Quick Check:

   Purpose of np.einsum() = To perform efficient and readable array operations using Einstein summation notation [OK]

Hint: Remember: einsum is for smart, fast array math [OK]

Common Mistakes:

• Confusing einsum with random array creation
• Thinking einsum sorts or reshapes arrays
• Assuming einsum is only for dot products

2. Which of the following is the correct syntax to compute the matrix product of two 2D arrays A and B using np.einsum()?

easy

A. np.einsum('ij,jk->ik', A, B)

B. np.einsum('ik,jk->ij', A, B)

C. np.einsum('ij,ik->jk', A, B)

D. np.einsum('ji,kj->ik', A, B)

Solution

1. Step 1: Recall matrix multiplication indices

   Matrix multiplication of A (shape ixj) and B (shape jxk) results in shape ixk, summing over j.

2. Step 2: Match einsum subscripts

   np.einsum('ij,jk->ik', A, B) uses 'ij,jk->ik', which correctly sums over j and produces ixk output.

3. Final Answer:

   np.einsum('ij,jk->ik', A, B) -> Option A

4. Quick Check:

   Matrix product einsum = np.einsum('ij,jk->ik', A, B) [OK]

Hint: Use 'ij,jk->ik' for matrix multiplication [OK]

Common Mistakes:

• Mixing up indices order in the subscript string
• Using wrong output indices after ->
• Confusing which axis to sum over

3. Given the code:

import numpy as np
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
result = np.einsum('ij,ij->i', A, B)
print(result)

What is the output?

medium

A. [19 43]

B. [17 43]

C. [23 53]

D. [17 53]

Solution

1. Step 1: Understand the einsum subscripts 'ij,ij->i'

   This sums element-wise multiplication over the second index j, keeping the first index i.

2. Step 2: Calculate for each i

   For i=0: (1*5)+(2*6)=5+12=17; For i=1: (3*7)+(4*8)=21+32=53.

3. Final Answer:

   [17 53] -> Option D

4. Quick Check:

   Element-wise sum over j = [17 53] [OK]

Hint: Sum over repeated indices, keep others [OK]

Common Mistakes:

• Summing over wrong axis
• Confusing output indices in einsum
• Calculating element-wise product incorrectly

4. Identify the error in this code snippet:

import numpy as np
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
result = np.einsum('ij,jk->ik', A, B)
print(result)

and choose the correct fix.

medium

A. Ensure B has shape (2, 2) and use 'ij,jk->ik' as is

B. No error; code runs correctly

C. Swap arrays to np.einsum('jk,ij->ik', B, A)

D. Change subscript to 'ij,ik->jk' to match shapes

Solution

1. Step 1: Check array shapes and einsum subscripts

   A and B both have shape (2,2). The subscript 'ij,jk->ik' expects B's first dimension (j) to match A's second dimension (j), which it does.

2. Step 2: Verify if code runs

   Shapes match perfectly for matrix multiplication. No error; the code runs correctly.

3. Final Answer:

   No error; code runs correctly -> Option B

4. Quick Check:

   No error; code runs correctly [OK]

Hint: Check array shapes match einsum indices [OK]

Common Mistakes:

• Assuming code errors without checking shapes
• Changing subscripts incorrectly
• Swapping arrays unnecessarily

5. You have a 3D array data with shape (10, 5, 4) representing 10 samples, 5 features, and 4 time points. Using np.einsum(), how do you compute the sum of products of features and time points for each sample, resulting in a 1D array of length 10?

hard

A. np.einsum('ijk,ijk->k', data, data)

B. np.einsum('ijk,ijk->j', data, data)

C. np.einsum('ijk,ijk->i', data, data)

D. np.einsum('ijk,ijk->ij', data, data)

Solution

1. Step 1: Understand the dimensions and desired output

   Data shape is (samples=10, features=5, time=4). We want to sum over features and time, keeping samples dimension.

2. Step 2: Use einsum to sum over j and k, keep i

   Subscript 'ijk,ijk->i' multiplies element-wise and sums over j and k, resulting in length 10 array.

3. Final Answer:

   np.einsum('ijk,ijk->i', data, data) -> Option C

4. Quick Check:

   Sum over features and time = np.einsum('ijk,ijk->i', data, data) [OK]

Hint: Sum over repeated indices, keep sample index [OK]

Common Mistakes:

• Summing over wrong axes
• Confusing output indices in einsum
• Expecting 2D output instead of 1D