1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
{% set name = "opt_einsum" %}
{% set version = "2.0.0" %}
package:
name: {{ name|lower }}
version: {{ version }}
source:
path: ../..
build:
number: 0
script: python setup.py install --single-version-externally-managed --record record.txt
requirements:
build:
- python
- setuptools
run:
- python
- numpy
test:
requires:
- python
- pytest
commands:
- py.test --pyargs opt_einsum
about:
home: http://github.com/dgasmith/opt_einsum
license: MIT
license_family: MIT
license_file: LICENSE
summary: 'A contraction optimizer for the NumPy Einsum function.'
description: >
Einsum is a very powerful function for contracting tensors of arbitrary dimension and index.
However, it is only optimized to contract two terms at a time resulting in non-optimal scaling.
For example, let us examine the following index transformation:
`M_{pqrs} = C_{pi} C_{qj} I_{ijkl} C_{rk} C_{sl}`
We can then develop two separate implementations that produce the same result:
```python
N = 10
C = np.random.rand(N, N)
I = np.random.rand(N, N, N, N)
def naive(I, C):
# N^8 scaling
return np.einsum('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C)
def optimized(I, C):
# N^5 scaling
K = np.einsum('pi,ijkl->pjkl', C, I)
K = np.einsum('qj,pjkl->pqkl', C, K)
K = np.einsum('rk,pqkl->pqrl', C, K)
K = np.einsum('sl,pqrl->pqrs', C, K)
return K
```
The einsum function does not consider building intermediate arrays; therefore, helping einsum out by building these intermediate arrays can result in a considerable cost savings even for small N (N=10):
```python
np.allclose(naive(I, C), optimized(I, C))
True
%timeit naive(I, C)
1 loops, best of 3: 934 ms per loop
%timeit optimized(I, C)
1000 loops, best of 3: 527 µs per loop
```
A 2000 fold speed up for 4 extra lines of code!
This contraction can be further complicated by considering that the shape of the C matrices need not be the same, in this case the ordering in which the indices are transformed matters greatly.
Logic can be built that optimizes the ordering; however, this is a lot of time and effort for a single expression.
The opt_einsum package is a drop in replacement for the np.einsum function and can handle all of this logic for you:
```python
from opt_einsum import contract
%timeit contract('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C)
1000 loops, best of 3: 324 µs per loop
```
The above will automatically find the optimal contraction order, in this case identical to that of the optimized function above, and compute the products for you. In this case, it even uses `np.dot` under the hood to exploit any vendor BLAS functionality that your NumPy build has!
dev_url: https://github.com/dgasmith/opt_einsum
extra:
recipe-maintainers:
- dgasmith
- loriab
|
