opt_einsum_optimized_tensor_contraction_path_quickstart.py

1import numpy as np
2from opt_einsum import contract
3
4# Define the dimensions
5N = 10
6C, D, E, F, G, H = N, N, N, N, N, N
7
8# Create random tensors
9a = np.random.rand(C, D)
10b = np.random.rand(D, E)
11c = np.random.rand(E, F)
12d = np.random.rand(F, G)
13e = np.random.rand(G, H)
14
15# Perform the optimized contraction
16# This will automatically find a path that minimizes the number of operations
17result = contract('cd,de,ef,fg,gh->ch', a, b, c, d, e)
18
19print(f"Result shape: {result.shape}")