einops-0.4.1.dist-info/INSTALLER ADDED
@@ -0,0 +1 @@
 
 
1
+ pip
einops-0.4.1.dist-info/LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2018 Alex Rogozhnikov
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
einops-0.4.1.dist-info/METADATA ADDED
@@ -0,0 +1,309 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Metadata-Version: 2.1
2
+ Name: einops
3
+ Version: 0.4.1
4
+ Summary: A new flavour of deep learning operations
5
+ Home-page: https://github.com/arogozhnikov/einops
6
+ Author: Alex Rogozhnikov
7
+ License: UNKNOWN
8
+ Keywords: deep learning,neural networks,tensor manipulation,machine learning,scientific computations,einops
9
+ Platform: UNKNOWN
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: Programming Language :: Python :: 3
12
+ Description-Content-Type: text/markdown
13
+ License-File: LICENSE
14
+
15
+ <!-- this link magically rendered as video, unfortunately not in docs -->
16
+
17
+ <!-- https://user-images.githubusercontent.com/6318811/116849688-0ca41c00-aba4-11eb-8ccf-74744f6cbc23.mp4 -->
18
+
19
+ <a href='http://arogozhnikov.github.io/images/einops/einops_video.mp4' >
20
+ <div align="center">
21
+ <img src="http://arogozhnikov.github.io/images/einops/einops_video.gif" alt="einops package examples" />
22
+ <br>
23
+ <small><a href='http://arogozhnikov.github.io/images/einops/einops_video.mp4'>This video in high quality (mp4)</a></small>
24
+ <br><br>
25
+ </div>
26
+ </a>
27
+
28
+ # einops
29
+ [![Run tests](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml/badge.svg)](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml)
30
+ [![PyPI version](https://badge.fury.io/py/einops.svg)](https://badge.fury.io/py/einops)
31
+ [![Documentation](https://img.shields.io/badge/documentation-link-blue.svg)](https://einops.rocks/)
32
+ ![Supported python versions](https://raw.githubusercontent.com/arogozhnikov/einops/master/docs/resources/python_badge.svg)
33
+
34
+
35
+ Flexible and powerful tensor operations for readable and reliable code.
36
+ Supports numpy, pytorch, tensorflow, jax, and [others](#supported-frameworks).
37
+
38
+ ## Recent updates:
39
+
40
+ - torch.jit.script is supported for pytorch layers
41
+ - powerful EinMix added to einops. [Einmix tutorial notebook](https://github.com/arogozhnikov/einops/blob/master/docs/3-einmix-layer.ipynb)
42
+
43
+ <!--<div align="center">
44
+ <img src="http://arogozhnikov.github.io/images/einops/einops_logo_350x350.png"
45
+ alt="einops package logo" width="250" height="250" />
46
+ <br><br>
47
+ </div> -->
48
+
49
+ ## Tweets
50
+
51
+ > In case you need convincing arguments for setting aside time to learn about einsum and einops...
52
+ [Tim Rocktäschel, FAIR](https://twitter.com/_rockt/status/1230818967205425152)
53
+
54
+ > Writing better code with PyTorch and einops 👌
55
+ [Andrej Karpathy, AI at Tesla](https://twitter.com/karpathy/status/1290826075916779520)
56
+
57
+ > Slowly but surely, einops is seeping in to every nook and cranny of my code. If you find yourself shuffling around bazillion dimensional tensors, this might change your life
58
+ [Nasim Rahaman, MILA (Montreal)](https://twitter.com/nasim_rahaman/status/1216022614755463169)
59
+
60
+ [More testimonials](https://einops.rocks/pages/testimonials/)
61
+
62
+ ## Contents
63
+
64
+ - [Installation](#Installation)
65
+ - [Documentation](https://einops.rocks/)
66
+ - [Tutorial](#Tutorials)
67
+ - [API micro-reference](#API)
68
+ - [Why using einops](#Why-using-einops-notation)
69
+ - [Supported frameworks](#Supported-frameworks)
70
+ - [Contributing](#Contributing)
71
+ - [Repository](https://github.com/arogozhnikov/einops) and [discussions](https://github.com/arogozhnikov/einops/discussions)
72
+
73
+ ## Installation <a name="Installation"></a>
74
+
75
+ Plain and simple:
76
+ ```bash
77
+ pip install einops
78
+ ```
79
+
80
+ <!--
81
+ `einops` has no mandatory dependencies (code examples also require jupyter, pillow + backends).
82
+ To obtain the latest github version
83
+
84
+ ```bash
85
+ pip install https://github.com/arogozhnikov/einops/archive/master.zip
86
+ ```
87
+ -->
88
+
89
+ ## Tutorials <a name="Tutorials"></a>
90
+
91
+ Tutorials are the most convenient way to see `einops` in action
92
+
93
+ - part 1: [einops fundamentals](https://github.com/arogozhnikov/einops/blob/master/docs/1-einops-basics.ipynb)
94
+ - part 2: [einops for deep learning](https://github.com/arogozhnikov/einops/blob/master/docs/2-einops-for-deep-learning.ipynb)
95
+ - part 3: [improve pytorch code with einops](https://arogozhnikov.github.io/einops/pytorch-examples.html)
96
+
97
+
98
+ ## API <a name="API"></a>
99
+
100
+ `einops` has a minimalistic yet powerful API.
101
+
102
+ Three operations provided ([einops tutorial](https://github.com/arogozhnikov/einops/blob/master/docs/)
103
+ shows those cover stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate, view and numerous reductions)
104
+
105
+ ```python
106
+ from einops import rearrange, reduce, repeat
107
+ # rearrange elements according to the pattern
108
+ output_tensor = rearrange(input_tensor, 't b c -> b c t')
109
+ # combine rearrangement and reduction
110
+ output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2)
111
+ # copy along a new axis
112
+ output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)
113
+ ```
114
+ And two corresponding layers (`einops` keeps a separate version for each framework) with the same API.
115
+
116
+ ```python
117
+ from einops.layers.chainer import Rearrange, Reduce
118
+ from einops.layers.gluon import Rearrange, Reduce
119
+ from einops.layers.keras import Rearrange, Reduce
120
+ from einops.layers.torch import Rearrange, Reduce
121
+ from einops.layers.tensorflow import Rearrange, Reduce
122
+ ```
123
+
124
+ Layers behave similarly to operations and have the same parameters
125
+ (with the exception of the first argument, which is passed during call)
126
+
127
+ ```python
128
+ layer = Rearrange(pattern, **axes_lengths)
129
+ layer = Reduce(pattern, reduction, **axes_lengths)
130
+
131
+ # apply created layer to a tensor / variable
132
+ x = layer(x)
133
+ ```
134
+
135
+ Example of using layers within a model:
136
+ ```python
137
+ # example given for pytorch, but code in other frameworks is almost identical
138
+ from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
139
+ from einops.layers.torch import Rearrange
140
+
141
+ model = Sequential(
142
+ Conv2d(3, 6, kernel_size=5),
143
+ MaxPool2d(kernel_size=2),
144
+ Conv2d(6, 16, kernel_size=5),
145
+ MaxPool2d(kernel_size=2),
146
+ # flattening
147
+ Rearrange('b c h w -> b (c h w)'),
148
+ Linear(16*5*5, 120),
149
+ ReLU(),
150
+ Linear(120, 10),
151
+ )
152
+ ```
153
+
154
+ <!---
155
+ Additionally two auxiliary functions provided
156
+ ```python
157
+ from einops import asnumpy, parse_shape
158
+ # einops.asnumpy converts tensors of imperative frameworks to numpy
159
+ numpy_tensor = asnumpy(input_tensor)
160
+ # einops.parse_shape gives a shape of axes of interest
161
+ parse_shape(input_tensor, 'batch _ h w') # e.g {'batch': 64, 'h': 128, 'w': 160}
162
+ ```
163
+ -->
164
+
165
+ ## Naming <a name="Naming"></a>
166
+
167
+ `einops` stands for Einstein-Inspired Notation for operations
168
+ (though "Einstein operations" is more attractive and easier to remember).
169
+
170
+ Notation was loosely inspired by Einstein summation (in particular by `numpy.einsum` operation).
171
+
172
+ ## Why use `einops` notation?! <a name="Why-using-einops-notation"></a>
173
+
174
+
175
+ ### Semantic information (being verbose in expectations)
176
+
177
+ ```python
178
+ y = x.view(x.shape[0], -1)
179
+ y = rearrange(x, 'b c h w -> b (c h w)')
180
+ ```
181
+ While these two lines are doing the same job in *some* context,
182
+ the second one provides information about the input and output.
183
+ In other words, `einops` focuses on interface: *what is the input and output*, not *how* the output is computed.
184
+
185
+ The next operation looks similar:
186
+
187
+ ```python
188
+ y = rearrange(x, 'time c h w -> time (c h w)')
189
+ ```
190
+ but it gives the reader a hint:
191
+ this is not an independent batch of images we are processing,
192
+ but rather a sequence (video).
193
+
194
+ Semantic information makes the code easier to read and maintain.
195
+
196
+ ### Convenient checks
197
+
198
+ Reconsider the same example:
199
+
200
+ ```python
201
+ y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
202
+ y = rearrange(x, 'b c h w -> b (c h w)')
203
+ ```
204
+ The second line checks that the input has four dimensions,
205
+ but you can also specify particular dimensions.
206
+ That's opposed to just writing comments about shapes since
207
+ [comments don't work and don't prevent mistakes](https://medium.freecodecamp.org/code-comments-the-good-the-bad-and-the-ugly-be9cc65fbf83)
208
+ as we know
209
+ ```python
210
+ y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
211
+ y = rearrange(x, 'b c h w -> b (c h w)', c=256, h=19, w=19)
212
+ ```
213
+
214
+ ### Result is strictly determined
215
+
216
+ Below we have at least two ways to define the depth-to-space operation
217
+ ```python
218
+ # depth-to-space
219
+ rearrange(x, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=2, w2=2)
220
+ rearrange(x, 'b c (h h2) (w w2) -> b (h2 w2 c) h w', h2=2, w2=2)
221
+ ```
222
+ There are at least four more ways to do it. Which one is used by the framework?
223
+
224
+ These details are ignored, since *usually* it makes no difference,
225
+ but it can make a big difference (e.g. if you use grouped convolutions in the next stage),
226
+ and you'd like to specify this in your code.
227
+
228
+
229
+ ### Uniformity
230
+
231
+ ```python
232
+ reduce(x, 'b c (x dx) -> b c x', 'max', dx=2)
233
+ reduce(x, 'b c (x dx) (y dy) -> b c x y', 'max', dx=2, dy=3)
234
+ reduce(x, 'b c (x dx) (y dy) (z dz) -> b c x y z', 'max', dx=2, dy=3, dz=4)
235
+ ```
236
+ These examples demonstrated that we don't use separate operations for 1d/2d/3d pooling,
237
+ those are all defined in a uniform way.
238
+
239
+ Space-to-depth and depth-to space are defined in many frameworks but how about width-to-height? Here you go:
240
+
241
+ ```python
242
+ rearrange(x, 'b c h (w w2) -> b c (h w2) w', w2=2)
243
+ ```
244
+
245
+ ### Framework independent behavior
246
+
247
+ Even simple functions are defined differently by different frameworks
248
+
249
+ ```python
250
+ y = x.flatten() # or flatten(x)
251
+ ```
252
+
253
+ Suppose `x`'s shape was `(3, 4, 5)`, then `y` has shape ...
254
+
255
+ - numpy, cupy, chainer, pytorch: `(60,)`
256
+ - keras, tensorflow.layers, mxnet and gluon: `(3, 20)`
257
+
258
+ `einops` works the same way in all frameworks.
259
+
260
+ ### Independence of framework terminology
261
+
262
+ Example: `tile` vs `repeat` causes lots of confusion. To copy image along width:
263
+ ```python
264
+ np.tile(image, (1, 2)) # in numpy
265
+ image.repeat(1, 2) # pytorch's repeat ~ numpy's tile
266
+ ```
267
+
268
+ With einops you don't need to decipher which axis was repeated:
269
+ ```python
270
+ repeat(image, 'h w -> h (tile w)', tile=2) # in numpy
271
+ repeat(image, 'h w -> h (tile w)', tile=2) # in pytorch
272
+ repeat(image, 'h w -> h (tile w)', tile=2) # in tf
273
+ repeat(image, 'h w -> h (tile w)', tile=2) # in jax
274
+ repeat(image, 'h w -> h (tile w)', tile=2) # in mxnet
275
+ ... (etc.)
276
+ ```
277
+
278
+ Testimonials provide user's perspective on the same question.
279
+
280
+ ## Supported frameworks <a name="Supported-frameworks"></a>
281
+
282
+ Einops works with ...
283
+
284
+ - [numpy](http://www.numpy.org/)
285
+ - [pytorch](https://pytorch.org/)
286
+ - [tensorflow](https://www.tensorflow.org/)
287
+ - [jax](https://github.com/google/jax)
288
+ - [cupy](https://cupy.chainer.org/)
289
+ - [chainer](https://chainer.org/)
290
+ - [gluon](https://gluon.mxnet.io/)
291
+ - [tf.keras](https://www.tensorflow.org/guide/keras)
292
+ - [mxnet](https://mxnet.apache.org/) (experimental)
293
+
294
+
295
+ ## Contributing <a name="Contributing"></a>
296
+
297
+ Best ways to contribute are
298
+
299
+ - spread the word about `einops`
300
+ - if you like explaining things, more tutorials/tear-downs of implementations is welcome
301
+ - tutorials in other languages are very welcome
302
+ - do you have project/code example to share? Let me know in github discussions
303
+ - use `einops` in your papers!
304
+
305
+ ## Supported python versions
306
+
307
+ `einops` works with python 3.6 or later.
308
+
309
+
einops-0.4.1.dist-info/RECORD ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ einops-0.4.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
2
+ einops-0.4.1.dist-info/LICENSE,sha256=MNmENkKW9R_67K1LAe4SfpUlDFBokY1LZvyWIGcj5DQ,1073
3
+ einops-0.4.1.dist-info/METADATA,sha256=UdOBa4tijnwPJI48dGASJt4-czHTJ4LLiY4dfdRXffI,10737
4
+ einops-0.4.1.dist-info/RECORD,,
5
+ einops-0.4.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
6
+ einops-0.4.1.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
7
+ einops-0.4.1.dist-info/top_level.txt,sha256=zh9ckJ4QUP-fUBSO5-UKAcNKvC_lzMGMYS6_nnoT4Tc,7
8
+ einops/__init__.py,sha256=8uWtV9MPDvreSsu6VG0E-j2TlxrckMGW2Zy5fCRFu6I,297
9
+ einops/__pycache__/__init__.cpython-310.pyc,,
10
+ einops/__pycache__/_backends.cpython-310.pyc,,
11
+ einops/__pycache__/_torch_specific.cpython-310.pyc,,
12
+ einops/__pycache__/einops.cpython-310.pyc,,
13
+ einops/__pycache__/parsing.cpython-310.pyc,,
14
+ einops/_backends.py,sha256=xyh2XysbubzGMqyRZPu8ld5-1bob693ESp6vkZR4gV8,17132
15
+ einops/_torch_specific.py,sha256=-VuZFXozi6GPtKQFWvb7BhzswyOfjByXGG0GLvvhDXg,2804
16
+ einops/einops.py,sha256=LkY-JdbOUqf_GtuiskedET4aaSs0yJKfVaLoBwS6UP8,27879
17
+ einops/layers/__init__.py,sha256=JHwHQUP5sBIYhSwRrjhZYxGIdw8-UTEWUPbeEduBuBY,2824
18
+ einops/layers/__pycache__/__init__.cpython-310.pyc,,
19
+ einops/layers/__pycache__/_einmix.cpython-310.pyc,,
20
+ einops/layers/__pycache__/chainer.cpython-310.pyc,,
21
+ einops/layers/__pycache__/gluon.cpython-310.pyc,,
22
+ einops/layers/__pycache__/keras.cpython-310.pyc,,
23
+ einops/layers/__pycache__/tensorflow.cpython-310.pyc,,
24
+ einops/layers/__pycache__/torch.cpython-310.pyc,,
25
+ einops/layers/_einmix.py,sha256=k1Wt5z7KmJF9nj345ZhUXeRBcV1D2bNkz35yF82zB_E,8249
26
+ einops/layers/chainer.py,sha256=VisqqyZiEpDl7NdCSjVSa4u7aXgZuNpA0hglkfGydiM,1927
27
+ einops/layers/gluon.py,sha256=Ll85s1OWKqRAhSwFS33jQwbTicD1MnhrH4lbnlqvoPU,2101
28
+ einops/layers/keras.py,sha256=RTsR-aim1Sco5VXI2W1Qs639hJRJ0hWIilTZCs3Ftn4,212
29
+ einops/layers/tensorflow.py,sha256=xNsVaKIMoB2kZeSeFUKXq29LWz-Fppt2K2aRln5s0-Y,3269
30
+ einops/layers/torch.py,sha256=IOdwPR2uL_ZFuzWthGz6p-8af1zg801UmjB8uTBA5HY,2379
31
+ einops/parsing.py,sha256=75hvgp6iWvvLUe67IaQujmox1tjvF9ZsBMaXQYnQmqU,6637
einops-0.4.1.dist-info/REQUESTED ADDED
File without changes
einops-0.4.1.dist-info/WHEEL ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ Wheel-Version: 1.0
2
+ Generator: bdist_wheel (0.37.1)
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
5
+
einops-0.4.1.dist-info/top_level.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ einops
einops/__init__.py ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ from .rearrange import rearrange
2
+
3
+ __all__ = ["rearrange"]
einops/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (248 Bytes). View file
 
einops/__pycache__/_parsing.cpython-310.pyc ADDED
Binary file (9.95 kB). View file
 
einops/__pycache__/rearrange.cpython-310.pyc ADDED
Binary file (7.22 kB). View file
 
einops/_parsing.py ADDED
@@ -0,0 +1,302 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Adapted from https://github.com/arogozhnikov/einops/blob/36c7bb16e57d6e57f8f3050f9e07abdf3f00469f/einops/parsing.py.
2
+
3
+ MIT License
4
+
5
+ Copyright (c) 2018 Alex Rogozhnikov
6
+
7
+ Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ of this software and associated documentation files (the "Software"), to deal
9
+ in the Software without restriction, including without limitation the rights
10
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ copies of the Software, and to permit persons to whom the Software is
12
+ furnished to do so, subject to the following conditions:
13
+
14
+ The above copyright notice and this permission notice shall be included in all
15
+ copies or substantial portions of the Software.
16
+
17
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ SOFTWARE.
24
+ """
25
+ from __future__ import annotations
26
+
27
+ import keyword
28
+ import warnings
29
+ from typing import Collection, List, Mapping, Optional, Set, Tuple, Union
30
+
31
+ _ellipsis: str = "…" # NB, this is a single unicode symbol. String is used as it is not a list, but can be iterated
32
+
33
+
34
+ class AnonymousAxis:
35
+ """Used by `ParsedExpression` to represent an axis with a size (> 1), but no associated identifier.
36
+
37
+ Note: Different instances of this class are not equal to each other, even if they have the same value.
38
+ """
39
+
40
+ def __init__(self, value: str) -> None:
41
+ self.value = int(value)
42
+ if self.value < 1:
43
+ raise ValueError(
44
+ f"Anonymous axis should have positive length, not {self.value}"
45
+ )
46
+
47
+ def __repr__(self) -> str:
48
+ return f"{self.value}-axis"
49
+
50
+
51
+ class ParsedExpression:
52
+ """Structure containing information about one side of an `einops`-style pattern (e.g. 'b c (h w)')."""
53
+
54
+ def __init__(
55
+ self,
56
+ expression: str,
57
+ *,
58
+ allow_underscore: bool = False,
59
+ allow_duplicates: bool = False,
60
+ ) -> None:
61
+ """Parse the expression and store relevant metadata.
62
+
63
+ Args:
64
+ expression (str): the `einops`-pattern to parse
65
+ allow_underscore (bool): whether to allow axis identifier names to begin with an underscore
66
+ allow_duplicates (bool): whether to allow an identifier to appear more than once in the expression
67
+ """
68
+ self.has_ellipsis: bool = False
69
+ self.has_ellipsis_parenthesized: Optional[bool] = None
70
+ self.identifiers: Set[Union[str, AnonymousAxis]] = set()
71
+ # that's axes like 2, 3, 4 or 5. Axes with size 1 are exceptional and replaced with empty composition
72
+ self.has_non_unitary_anonymous_axes: bool = False
73
+ # composition keeps structure of composite axes, see how different corner cases are handled in tests
74
+ self.composition: List[Union[List[Union[str, AnonymousAxis]], str]] = []
75
+ if "." in expression:
76
+ if "..." not in expression:
77
+ raise ValueError(
78
+ "Expression may contain dots only inside ellipsis (...)"
79
+ )
80
+ if str.count(expression, "...") != 1 or str.count(expression, ".") != 3:
81
+ raise ValueError(
82
+ "Expression may contain dots only inside ellipsis (...); only one ellipsis for tensor "
83
+ )
84
+ expression = expression.replace("...", _ellipsis)
85
+ self.has_ellipsis = True
86
+
87
+ bracket_group: Optional[List[Union[str, AnonymousAxis]]] = None
88
+
89
+ def add_axis_name(x: str) -> None:
90
+ if x in self.identifiers:
91
+ if not (allow_underscore and x == "_") and not allow_duplicates:
92
+ raise ValueError(
93
+ f"Indexing expression contains duplicate dimension '{x}'"
94
+ )
95
+ if x == _ellipsis:
96
+ self.identifiers.add(_ellipsis)
97
+ if bracket_group is None:
98
+ self.composition.append(_ellipsis)
99
+ self.has_ellipsis_parenthesized = False
100
+ else:
101
+ bracket_group.append(_ellipsis)
102
+ self.has_ellipsis_parenthesized = True
103
+ else:
104
+ is_number = str.isdecimal(x)
105
+ if is_number and int(x) == 1:
106
+ # handling the case of anonymous axis of length 1
107
+ if bracket_group is None:
108
+ self.composition.append([])
109
+ else:
110
+ pass # no need to think about 1s inside parenthesis
111
+ return
112
+ is_axis_name, reason = self.check_axis_name_return_reason(
113
+ x, allow_underscore=allow_underscore
114
+ )
115
+ if not (is_number or is_axis_name):
116
+ raise ValueError(f"Invalid axis identifier: {x}\n{reason}")
117
+ axis_name: Union[str, AnonymousAxis] = (
118
+ AnonymousAxis(x) if is_number else x
119
+ )
120
+ self.identifiers.add(axis_name)
121
+ if is_number:
122
+ self.has_non_unitary_anonymous_axes = True
123
+ if bracket_group is None:
124
+ self.composition.append([axis_name])
125
+ else:
126
+ bracket_group.append(axis_name)
127
+
128
+ current_identifier = None
129
+ for char in expression:
130
+ if char in "() ":
131
+ if current_identifier is not None:
132
+ add_axis_name(current_identifier)
133
+ current_identifier = None
134
+ if char == "(":
135
+ if bracket_group is not None:
136
+ raise ValueError(
137
+ "Axis composition is one-level (brackets inside brackets not allowed)"
138
+ )
139
+ bracket_group = []
140
+ elif char == ")":
141
+ if bracket_group is None:
142
+ raise ValueError("Brackets are not balanced")
143
+ self.composition.append(bracket_group)
144
+ bracket_group = None
145
+ elif str.isalnum(char) or char in ["_", _ellipsis]:
146
+ if current_identifier is None:
147
+ current_identifier = char
148
+ else:
149
+ current_identifier += char
150
+ else:
151
+ raise ValueError(f"Unknown character '{char}'")
152
+
153
+ if bracket_group is not None:
154
+ raise ValueError(f"Imbalanced parentheses in expression: '{expression}'")
155
+ if current_identifier is not None:
156
+ add_axis_name(current_identifier)
157
+
158
+ @staticmethod
159
+ def check_axis_name_return_reason(
160
+ name: str, allow_underscore: bool = False
161
+ ) -> Tuple[bool, str]:
162
+ """Check if the given axis name is valid, and a message explaining why if not.
163
+
164
+ Valid axes names are python identifiers except keywords, and should not start or end with an underscore.
165
+
166
+ Args:
167
+ name (str): the axis name to check
168
+ allow_underscore (bool): whether axis names are allowed to start with an underscore
169
+
170
+ Returns:
171
+ Tuple[bool, str]: whether the axis name is valid, a message explaining why if not
172
+ """
173
+ if not str.isidentifier(name):
174
+ return False, "not a valid python identifier"
175
+ elif name[0] == "_" or name[-1] == "_":
176
+ if name == "_" and allow_underscore:
177
+ return True, ""
178
+ return False, "axis name should should not start or end with underscore"
179
+ else:
180
+ if keyword.iskeyword(name):
181
+ warnings.warn(
182
+ f"It is discouraged to use axes names that are keywords: {name}",
183
+ RuntimeWarning,
184
+ )
185
+ if name in ["axis"]:
186
+ warnings.warn(
187
+ "It is discouraged to use 'axis' as an axis name and will raise an error in future",
188
+ FutureWarning,
189
+ )
190
+ return True, ""
191
+
192
+ @staticmethod
193
+ def check_axis_name(name: str) -> bool:
194
+ """Check if the name is a valid axis name.
195
+
196
+ Args:
197
+ name (str): the axis name to check
198
+
199
+ Returns:
200
+ bool: whether the axis name is valid
201
+ """
202
+ is_valid, _ = ParsedExpression.check_axis_name_return_reason(name)
203
+ return is_valid
204
+
205
+
206
+ def parse_pattern(
207
+ pattern: str, axes_lengths: Mapping[str, int]
208
+ ) -> Tuple[ParsedExpression, ParsedExpression]:
209
+ """Parse an `einops`-style pattern into a left-hand side and right-hand side `ParsedExpression` object.
210
+
211
+ Args:
212
+ pattern (str): the `einops`-style rearrangement pattern
213
+ axes_lengths (Mapping[str, int]): any additional length specifications for dimensions
214
+
215
+ Returns:
216
+ Tuple[ParsedExpression, ParsedExpression]: a tuple containing the left-hand side and right-hand side expressions
217
+ """
218
+ # adapted from einops.einops._prepare_transformation_recipe
219
+ # https://github.com/arogozhnikov/einops/blob/230ac1526c1f42c9e1f7373912c7f8047496df11/einops/einops.py
220
+ try:
221
+ left_str, right_str = pattern.split("->")
222
+ except ValueError:
223
+ raise ValueError("Pattern must contain a single '->' separator") from None
224
+
225
+ if _ellipsis in axes_lengths:
226
+ raise ValueError(f"'{_ellipsis}' is not an allowed axis identifier")
227
+
228
+ left = ParsedExpression(left_str)
229
+ right = ParsedExpression(right_str)
230
+
231
+ if not left.has_ellipsis and right.has_ellipsis:
232
+ raise ValueError(
233
+ f"Ellipsis found in right side, but not left side of a pattern {pattern}"
234
+ )
235
+ if left.has_ellipsis and left.has_ellipsis_parenthesized:
236
+ raise ValueError(
237
+ f"Ellipsis is parenthesis in the left side is not allowed: {pattern}"
238
+ )
239
+
240
+ return left, right
241
+
242
+
243
+ def validate_rearrange_expressions(
244
+ left: ParsedExpression, right: ParsedExpression, axes_lengths: Mapping[str, int]
245
+ ) -> None:
246
+ """Perform expression validations that are specific to the `rearrange` operation.
247
+
248
+ Args:
249
+ left (ParsedExpression): left-hand side expression
250
+ right (ParsedExpression): right-hand side expression
251
+ axes_lengths (Mapping[str, int]): any additional length specifications for dimensions
252
+ """
253
+ for length in axes_lengths.values():
254
+ if (length_type := type(length)) is not int:
255
+ raise TypeError(
256
+ f"rearrange axis lengths must be integers, got: {length_type}"
257
+ )
258
+
259
+ if left.has_non_unitary_anonymous_axes or right.has_non_unitary_anonymous_axes:
260
+ raise ValueError("rearrange only supports unnamed axes of size 1")
261
+
262
+ difference = set.symmetric_difference(left.identifiers, right.identifiers)
263
+ if len(difference) > 0:
264
+ raise ValueError(
265
+ f"Identifiers only on one side of rearrange expression (should be on both): {difference}"
266
+ )
267
+
268
+ unmatched_axes = axes_lengths.keys() - left.identifiers
269
+ if len(unmatched_axes) > 0:
270
+ raise ValueError(
271
+ f"Identifiers not found in rearrange expression: {unmatched_axes}"
272
+ )
273
+
274
+
275
+ def comma_separate(collection: Collection[Union[str, Collection[str]]]) -> str:
276
+ """Convert a collection of strings representing first class dims into a comma-separated string.
277
+
278
+ Args:
279
+ collection (Collection[Union[str, Collection[str]]]): the collection of strings to convert
280
+
281
+ Returns:
282
+ str: the comma-separated string
283
+
284
+ Examples:
285
+ >>> comma_separate(('d0',))
286
+ 'd0'
287
+
288
+ >>> comma_separate(('d0', 'd1', 'd2', 'd3'))
289
+ 'd0, d1, d2, d3'
290
+
291
+ >>> comma_separate([('d1', 'd4')])
292
+ '(d1, d4)'
293
+
294
+ >>> comma_separate([('d0',), (), ('d1',), ('d2',), ('d3', 'd4')])
295
+ '(d0,), (), (d1,), (d2,), (d3, d4)'
296
+ """
297
+ return ", ".join(
298
+ item
299
+ if isinstance(item, str)
300
+ else f"({comma_separate(item)}{',' if len(item) == 1 else ''})"
301
+ for item in collection
302
+ )
einops/rearrange.py ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import functools
4
+ from typing import Callable, Dict, List, Sequence, Tuple, Union
5
+
6
+ import torch
7
+
8
+ from functorch._C import dim as _C
9
+ from ._parsing import (
10
+ _ellipsis,
11
+ AnonymousAxis,
12
+ comma_separate,
13
+ parse_pattern,
14
+ validate_rearrange_expressions,
15
+ )
16
+
17
+ __all__ = ["rearrange"]
18
+
19
+ dims = _C.dims
20
+
21
+
22
+ @functools.lru_cache(256)
23
+ def _create_rearrange_callable(
24
+ tensor_ndim: int, pattern: str, **axes_lengths: int
25
+ ) -> Callable[[torch.Tensor], torch.Tensor]:
26
+ r"""Translate an `einops`-style pattern into a callable that performs the rearrange using first-class dimensions.
27
+
28
+ Since the an equivalent result is computed for tensors with the same number of dimensions, with the same pattern and
29
+ specified axes lengths, this function can be memoized.
30
+
31
+ Args:
32
+ tensor_ndim (int): the number of dimensions in the tensor to rearrange
33
+ pattern (str): the `einops`-style rearrangement pattern
34
+ axes_lengths (int): any additional length specifications for dimensions
35
+
36
+ Returns:
37
+ Callable[[torch.Tensor], torch.Tensor]: a callable that performs the rearrangement
38
+ """
39
+ left, right = parse_pattern(pattern, axes_lengths)
40
+ validate_rearrange_expressions(left, right, axes_lengths)
41
+
42
+ n_anon_dims = sum(not dim for dim in left.composition)
43
+ if left.has_ellipsis:
44
+ n_ellipsis_dims = tensor_ndim - (len(left.composition) - 1)
45
+ n_named_dims = len(left.identifiers) - 1
46
+
47
+ if (pattern_ndim := n_anon_dims + n_named_dims) > tensor_ndim:
48
+ raise ValueError(
49
+ f"Number of dimensions in pattern ({pattern_ndim}) must be less than or equal to the number of "
50
+ f"dimensions in the tensor ({tensor_ndim})"
51
+ )
52
+ else:
53
+ n_ellipsis_dims = 0
54
+ n_named_dims = len(left.identifiers)
55
+
56
+ if (pattern_ndim := len(left.composition)) != tensor_ndim:
57
+ raise ValueError(
58
+ f"Number of dimensions in pattern ({pattern_ndim}) must be equal to the number of dimensions in "
59
+ f"the tensor ({tensor_ndim})"
60
+ )
61
+ n_dims = n_named_dims + n_ellipsis_dims + n_anon_dims
62
+
63
+ if n_dims == 0:
64
+ # an identity rearrangement on a 0-dimension tensor
65
+ return lambda tensor: tensor
66
+
67
+ first_class_dims: Tuple[str, ...] = tuple(f"d{i}" for i in range(n_dims))
68
+ identifier_dim_map: Dict[Union[str, AnonymousAxis], Tuple[str, ...]] = {}
69
+ anon_axes: List[AnonymousAxis] = []
70
+
71
+ # map the left-hand side identifiers to strings representing first class dims
72
+ dims_i = 0
73
+ for dimension in left.composition:
74
+ if isinstance(dimension, list):
75
+ for identifier in dimension:
76
+ # non-unitary anon axes are not allowed in rearrange & unitary anon axes are represented as empty lists
77
+ assert isinstance(identifier, str)
78
+ identifier_dim_map[identifier] = (first_class_dims[dims_i],)
79
+ dims_i += 1
80
+ if not dimension:
81
+ # unitary anonymous axis
82
+ anon_axis = AnonymousAxis("1")
83
+ identifier_dim_map[anon_axis] = (first_class_dims[dims_i],)
84
+ anon_axes.append(anon_axis)
85
+ dimension.append(anon_axis)
86
+ dims_i += 1
87
+ elif dimension == _ellipsis:
88
+ identifier = _ellipsis
89
+ identifier_dim_map[identifier] = tuple(
90
+ first_class_dims[dims_i + j] for j in range(n_ellipsis_dims)
91
+ )
92
+ dims_i += n_ellipsis_dims
93
+ else:
94
+ raise ValueError(f"Unexpected dimension: {dimension}")
95
+
96
+ def composition_to_dims(
97
+ composition: Sequence[Union[List[Union[str, AnonymousAxis]], str]]
98
+ ) -> List[Union[str, Tuple[str, ...]]]:
99
+ """Convert a `ParsedExpression.composition` into a `Tensor.__getitem__` index of strings representing first
100
+ class dims."""
101
+ dim_composition: List[Union[str, Tuple[str, ...]]] = []
102
+ for dimension in composition:
103
+ if isinstance(dimension, list):
104
+ dim_composition.append(
105
+ tuple(
106
+ dim
107
+ for identifier in dimension
108
+ for dim in identifier_dim_map[identifier]
109
+ )
110
+ )
111
+ elif dimension == _ellipsis:
112
+ dim_composition.extend(identifier_dim_map[_ellipsis])
113
+ else:
114
+ raise ValueError(f"Unexpected dimension: {dimension}")
115
+ return dim_composition
116
+
117
+ left_dims = composition_to_dims(left.composition)
118
+ right_dims = composition_to_dims(right.composition)
119
+ anon_dims = tuple(identifier_dim_map[axis][0] for axis in anon_axes)
120
+ specified_lengths = tuple(
121
+ (identifier_dim_map[axis][0], length) for axis, length in axes_lengths.items()
122
+ )
123
+
124
+ custom_rearrange_callable_name = "do_rearrange"
125
+ custom_rearrange_callable_code = (
126
+ (
127
+ f"def {custom_rearrange_callable_name}(tensor):\n"
128
+ f" {comma_separate(first_class_dims)} = dims({n_dims})\n"
129
+ )
130
+ + (
131
+ "".join(
132
+ f" {dim}.size = {length}\n" for (dim, length) in specified_lengths
133
+ )
134
+ if specified_lengths
135
+ else ""
136
+ )
137
+ + f" tensor = tensor[{comma_separate(left_dims)}].order({comma_separate(right_dims)})\n"
138
+ + (
139
+ f" return tensor.sum({comma_separate([anon_dims])}, keepdim=False)\n"
140
+ if anon_dims
141
+ else " return tensor\n"
142
+ )
143
+ )
144
+
145
+ exec(custom_rearrange_callable_code)
146
+ return locals()[custom_rearrange_callable_name]
147
+
148
+
149
+ def rearrange(
150
+ tensor: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor, ...]],
151
+ pattern: str,
152
+ **axes_lengths: int,
153
+ ) -> torch.Tensor:
154
+ r"""A native implementation of `einops.rearrange`, a reader-friendly smart element reordering for multidimensional
155
+ tensors. This operation includes functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze,
156
+ stack, concatenate and other operations.
157
+
158
+ See: https://einops.rocks/api/rearrange/
159
+
160
+ Args:
161
+ tensor (Tensor or sequence of Tensor): the tensor(s) to rearrange
162
+ pattern (str): the rearrangement pattern
163
+ axes_lengths (int): any additional length specifications for dimensions
164
+
165
+ Returns:
166
+ Tensor: the rearranged tensor
167
+
168
+ Examples:
169
+ >>> # suppose we have a set of 32 images in "h w c" format (height-width-channel)
170
+ >>> images = torch.randn((32, 30, 40, 3))
171
+
172
+ >>> # stack along first (batch) axis, output is a single array
173
+ >>> rearrange(images, 'b h w c -> b h w c').shape
174
+ torch.Size([32, 30, 40, 3])
175
+
176
+ >>> # concatenate images along height (vertical axis), 960 = 32 * 30
177
+ >>> rearrange(images, 'b h w c -> (b h) w c').shape
178
+ torch.Size([960, 40, 3])
179
+
180
+ >>> # concatenated images along horizontal axis, 1280 = 32 * 40
181
+ >>> rearrange(images, 'b h w c -> h (b w) c').shape
182
+ torch.Size([30, 1280, 3])
183
+
184
+ >>> # reordered axes to "b c h w" format for deep learning
185
+ >>> rearrange(images, 'b h w c -> b c h w').shape
186
+ torch.Size([32, 3, 30, 40])
187
+
188
+ >>> # flattened each image into a vector, 3600 = 30 * 40 * 3
189
+ >>> rearrange(images, 'b h w c -> b (c h w)').shape
190
+ torch.Size([32, 3600])
191
+
192
+ >>> # split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2
193
+ >>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
194
+ torch.Size([128, 15, 20, 3])
195
+
196
+ >>> # space-to-depth operation
197
+ >>> rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape
198
+ torch.Size([32, 15, 20, 12])
199
+ """
200
+ if not isinstance(tensor, torch.Tensor):
201
+ tensor = torch.stack(tensor)
202
+
203
+ rearrange_callable = _create_rearrange_callable(
204
+ tensor.ndim, pattern, **axes_lengths
205
+ )
206
+
207
+ return rearrange_callable(tensor)