METADATA 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394
  1. Metadata-Version: 2.4
  2. Name: einops
  3. Version: 0.8.2
  4. Summary: A new flavour of deep learning operations
  5. Project-URL: Homepage, https://github.com/arogozhnikov/einops
  6. Author: Alex Rogozhnikov
  7. License: MIT
  8. License-File: LICENSE
  9. Keywords: deep learning,einops,machine learning,neural networks,scientific computations,tensor manipulation
  10. Classifier: Intended Audience :: Science/Research
  11. Classifier: License :: OSI Approved :: MIT License
  12. Classifier: Programming Language :: Python :: 3
  13. Requires-Python: >=3.9
  14. Description-Content-Type: text/markdown
  15. <!--
  16. <a href='http://arogozhnikov.github.io/images/einops/einops_video.mp4' >
  17. <div align="center">
  18. <img src="http://arogozhnikov.github.io/images/einops/einops_video.gif" alt="einops package examples" />
  19. <br>
  20. <small><a href='http://arogozhnikov.github.io/images/einops/einops_video.mp4'>This video in high quality (mp4)</a></small>
  21. <br><br>
  22. </div>
  23. </a>
  24. -->
  25. <!-- this link magically rendered as video on github readme, unfortunately not in docs -->
  26. https://user-images.githubusercontent.com/6318811/177030658-66f0eb5d-e136-44d8-99c9-86ae298ead5b.mp4
  27. # einops
  28. [![Run tests](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml/badge.svg)](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml)
  29. [![PyPI version](https://badge.fury.io/py/einops.svg)](https://badge.fury.io/py/einops)
  30. [![Documentation](https://img.shields.io/badge/documentation-link-blue.svg)](https://einops.rocks/)
  31. ![Supported python versions](https://raw.githubusercontent.com/arogozhnikov/einops/main/docs/resources/python_badge.svg)
  32. Flexible and powerful tensor operations for readable and reliable code. <br />
  33. Supports numpy, pytorch, tensorflow, jax, and [others](#supported-frameworks).
  34. ## Recent updates:
  35. - 0.8.0: tinygrad backend added, small fixes
  36. - 0.7.0: no-hassle `torch.compile`, support of [array api standard](https://data-apis.org/array-api/latest/API_specification/index.html) and more
  37. - 10'000🎉: github reports that more than 10k project use einops
  38. - einops 0.6.1: paddle backend added
  39. - einops 0.6 introduces [packing and unpacking](https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb)
  40. - einops 0.5: einsum is now a part of einops
  41. - [Einops paper](https://openreview.net/pdf?id=oapKSVM2bcj) is accepted for oral presentation at ICLR 2022 (yes, it worth reading).
  42. Talk recordings are [available](https://iclr.cc/virtual/2022/oral/6603)
  43. <details markdown="1">
  44. <summary>Previous updates</summary>
  45. - flax and oneflow backend added
  46. - torch.jit.script is supported for pytorch layers
  47. - powerful EinMix added to einops. [Einmix tutorial notebook](https://github.com/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb)
  48. </details>
  49. <!--<div align="center">
  50. <img src="http://arogozhnikov.github.io/images/einops/einops_logo_350x350.png"
  51. alt="einops package logo" width="250" height="250" />
  52. <br><br>
  53. </div> -->
  54. ## Tweets
  55. > In case you need convincing arguments for setting aside time to learn about einsum and einops...
  56. [Tim Rocktäschel](https://twitter.com/_rockt/status/1230818967205425152)
  57. > Writing better code with PyTorch and einops 👌
  58. [Andrej Karpathy](https://twitter.com/karpathy/status/1290826075916779520)
  59. > 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
  60. [Nasim Rahaman](https://twitter.com/nasim_rahaman/status/1216022614755463169)
  61. [More testimonials](https://einops.rocks/pages/testimonials/)
  62. ## Contents
  63. - [Installation](#Installation)
  64. - [Documentation](https://einops.rocks/)
  65. - [Tutorial](#Tutorials)
  66. - [API micro-reference](#API)
  67. - [Why use einops](#Why-use-einops-notation)
  68. - [Supported frameworks](#Supported-frameworks)
  69. - [Citing](#Citing)
  70. - [Repository](https://github.com/arogozhnikov/einops) and [discussions](https://github.com/arogozhnikov/einops/discussions)
  71. ## Installation <a name="Installation"></a>
  72. Plain and simple:
  73. ```bash
  74. pip install einops
  75. ```
  76. (`uv pip install einops` works as well)
  77. ## Tutorials <a name="Tutorials"></a>
  78. Tutorials are the most convenient way to see `einops` in action
  79. - part 1: [einops fundamentals](https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb)
  80. - part 2: [einops for deep learning](https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb)
  81. - part 3: [packing and unpacking](https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb)
  82. - part 4: [improve pytorch code with einops](http://einops.rocks/pytorch-examples.html)
  83. Kapil Sachdeva recorded a small [intro to einops](https://www.youtube.com/watch?v=xGy75Pjsqzo).
  84. ## API <a name="API"></a>
  85. `einops` has a minimalistic yet powerful API.
  86. Three core operations provided ([einops tutorial](https://github.com/arogozhnikov/einops/blob/main/docs/)
  87. shows those cover stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate, view and numerous reductions)
  88. ```python
  89. from einops import rearrange, reduce, repeat
  90. # rearrange elements according to the pattern
  91. output_tensor = rearrange(input_tensor, 't b c -> b c t')
  92. # combine rearrangement and reduction
  93. output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2)
  94. # copy along a new axis
  95. output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)
  96. ```
  97. Later additions to the family are `pack` and `unpack` functions (better than stack/split/concatenate):
  98. ```python
  99. from einops import pack, unpack
  100. # pack and unpack allow reversibly 'packing' multiple tensors into one.
  101. # Packed tensors may be of different dimensionality:
  102. packed, ps = pack([class_token_bc, image_tokens_bhwc, text_tokens_btc], 'b * c')
  103. class_emb_bc, image_emb_bhwc, text_emb_btc = unpack(transformer(packed), ps, 'b * c')
  104. ```
  105. Finally, einops provides einsum with a support of multi-lettered names:
  106. ```python
  107. from einops import einsum, pack, unpack
  108. # einsum is like ... einsum, generic and flexible dot-product
  109. # but 1) axes can be multi-lettered 2) pattern goes last 3) works with multiple frameworks
  110. C = einsum(A, B, 'b t1 head c, b t2 head c -> b head t1 t2')
  111. ```
  112. ### EinMix
  113. `EinMix` is a generic linear layer, perfect for MLP Mixers and similar architectures.
  114. ### Layers
  115. Einops provides layers (`einops` keeps a separate version for each framework) that reflect corresponding functions
  116. ```python
  117. from einops.layers.torch import Rearrange, Reduce
  118. from einops.layers.tensorflow import Rearrange, Reduce
  119. from einops.layers.flax import Rearrange, Reduce
  120. from einops.layers.paddle import Rearrange, Reduce
  121. ```
  122. <details markdown="1">
  123. <summary>Example of using layers within a pytorch model</summary>
  124. Example given for pytorch, but code in other frameworks is almost identical
  125. ```python
  126. from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
  127. from einops.layers.torch import Rearrange
  128. model = Sequential(
  129. ...,
  130. Conv2d(6, 16, kernel_size=5),
  131. MaxPool2d(kernel_size=2),
  132. # flattening without need to write forward
  133. Rearrange('b c h w -> b (c h w)'),
  134. Linear(16*5*5, 120),
  135. ReLU(),
  136. Linear(120, 10),
  137. )
  138. ```
  139. No more flatten needed!
  140. Additionally, torch layers as those are script-able and compile-able.
  141. Operations [are torch.compile-able](https://github.com/arogozhnikov/einops/wiki/Using-torch.compile-with-einops),
  142. but not script-able due to limitations of torch.jit.script.
  143. </details>
  144. ## Naming <a name="Naming"></a>
  145. `einops` stands for Einstein-Inspired Notation for operations
  146. (though "Einstein operations" is more attractive and easier to remember).
  147. Notation was loosely inspired by Einstein summation (in particular by `numpy.einsum` operation).
  148. ## Why use `einops` notation?! <a name="Why-use-einops-notation"></a>
  149. ### Semantic information (being verbose in expectations)
  150. ```python
  151. y = x.view(x.shape[0], -1)
  152. y = rearrange(x, 'b c h w -> b (c h w)')
  153. ```
  154. While these two lines are doing the same job in *some* context,
  155. the second one provides information about the input and output.
  156. In other words, `einops` focuses on interface: *what is the input and output*, not *how* the output is computed.
  157. The next operation looks similar:
  158. ```python
  159. y = rearrange(x, 'time c h w -> time (c h w)')
  160. ```
  161. but it gives the reader a hint:
  162. this is not an independent batch of images we are processing,
  163. but rather a sequence (video).
  164. Semantic information makes the code easier to read and maintain.
  165. ### Convenient checks
  166. Reconsider the same example:
  167. ```python
  168. y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
  169. y = rearrange(x, 'b c h w -> b (c h w)')
  170. ```
  171. The second line checks that the input has four dimensions,
  172. but you can also specify particular dimensions.
  173. That's opposed to just writing comments about shapes since comments don't prevent mistakes,
  174. not tested, and without code review tend to be outdated
  175. ```python
  176. y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
  177. y = rearrange(x, 'b c h w -> b (c h w)', c=256, h=19, w=19)
  178. ```
  179. ### Result is strictly determined
  180. Below we have at least two ways to define the depth-to-space operation
  181. ```python
  182. # depth-to-space
  183. rearrange(x, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=2, w2=2)
  184. rearrange(x, 'b c (h h2) (w w2) -> b (h2 w2 c) h w', h2=2, w2=2)
  185. ```
  186. There are at least four more ways to do it. Which one is used by the framework?
  187. These details are ignored, since *usually* it makes no difference,
  188. but it can make a big difference (e.g. if you use grouped convolutions in the next stage),
  189. and you'd like to specify this in your code.
  190. ### Uniformity
  191. ```python
  192. reduce(x, 'b c (x dx) -> b c x', 'max', dx=2)
  193. reduce(x, 'b c (x dx) (y dy) -> b c x y', 'max', dx=2, dy=3)
  194. reduce(x, 'b c (x dx) (y dy) (z dz) -> b c x y z', 'max', dx=2, dy=3, dz=4)
  195. ```
  196. These examples demonstrated that we don't use separate operations for 1d/2d/3d pooling,
  197. those are all defined in a uniform way.
  198. Space-to-depth and depth-to space are defined in many frameworks but how about width-to-height? Here you go:
  199. ```python
  200. rearrange(x, 'b c h (w w2) -> b c (h w2) w', w2=2)
  201. ```
  202. ### Framework independent behavior
  203. Even simple functions are defined differently by different frameworks
  204. ```python
  205. y = x.flatten() # or flatten(x)
  206. ```
  207. Suppose `x`'s shape was `(3, 4, 5)`, then `y` has shape ...
  208. - numpy, pytorch, cupy, chainer, jax: `(60,)`
  209. - keras, tensorflow.layers, gluon: `(3, 20)`
  210. `einops` works the same way in all frameworks.
  211. ### Independence of framework terminology
  212. Example: `tile` vs `repeat` causes lots of confusion. To copy image along width:
  213. ```python
  214. np.tile(image, (1, 2)) # in numpy
  215. image.repeat(1, 2) # pytorch's repeat ~ numpy's tile
  216. ```
  217. With einops you don't need to decipher which axis was repeated:
  218. ```python
  219. repeat(image, 'h w -> h (tile w)', tile=2) # in numpy
  220. repeat(image, 'h w -> h (tile w)', tile=2) # in pytorch
  221. repeat(image, 'h w -> h (tile w)', tile=2) # in tf
  222. repeat(image, 'h w -> h (tile w)', tile=2) # in jax
  223. repeat(image, 'h w -> h (tile w)', tile=2) # in cupy
  224. ... (etc.)
  225. ```
  226. [Testimonials](https://einops.rocks/pages/testimonials/) provide users' perspective on the same question.
  227. ## Supported frameworks <a name="Supported-frameworks"></a>
  228. Einops works with ...
  229. - [numpy](http://www.numpy.org/)
  230. - [pytorch](https://pytorch.org/)
  231. - [tensorflow](https://www.tensorflow.org/)
  232. - [jax](https://github.com/google/jax)
  233. - [cupy](https://github.com/cupy/cupy)
  234. - [flax](https://github.com/google/flax) (community)
  235. - [paddle](https://github.com/PaddlePaddle/Paddle) (community)
  236. - [oneflow](https://github.com/Oneflow-Inc/oneflow) (community)
  237. - [tinygrad](https://github.com/tinygrad/tinygrad) (community)
  238. - [pytensor](https://github.com/pymc-devs/pytensor) (community)
  239. ```python
  240. from einops import rearrange
  241. => from einops.array_api import rearrange
  242. ```
  243. But actually it is even better: einops can be used with *any* framework that supports
  244. [Python array API standard](https://data-apis.org/array-api/latest/API_specification/index.html),
  245. to name a few:
  246. - numpy >= 2.0
  247. - [MLX](https://github.com/ml-explore/mlx) # yes, einops works with apple's framework
  248. - [pydata/sparse](https://github.com/pydata/sparse) >= 0.15 # and works with sparse tensors
  249. - [cubed](https://github.com/cubed-dev/cubed) # and with distributed tensors too
  250. - [quantco/ndonnx](https://github.com/Quantco/ndonnx)
  251. - jax
  252. - cupy
  253. - dask is supported via [array-api-compat](https://github.com/data-apis/array-api-compat)
  254. ## Development
  255. Devcontainer is provided, this environment can be used locally, or on your server,
  256. or within github codespaces.
  257. To start with devcontainers in vs code, clone repo, and click 'Reopen in Devcontainer'.
  258. Starting from einops 0.8.1, einops distributes tests as a part of package.
  259. ```bash
  260. # pip install einops pytest
  261. python -m einops.tests.run_tests numpy pytorch jax --pip-install
  262. ```
  263. `numpy pytorch jax` is an _example_, any subset of testable frameworks can be provided.
  264. Every framework is tested against numpy, so it is a requirement for tests.
  265. Specifying `--pip-install` will install requirements in current virtualenv,
  266. and should be omitted if dependencies are installed locally.
  267. To build/test docs:
  268. ```bash
  269. hatch run docs:serve # Serving on http://localhost:8000/
  270. ```
  271. ## Citing einops <a name="Citing"></a>
  272. Please use the following bibtex record
  273. ```text
  274. @inproceedings{
  275. rogozhnikov2022einops,
  276. title={Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation},
  277. author={Alex Rogozhnikov},
  278. booktitle={International Conference on Learning Representations},
  279. year={2022},
  280. url={https://openreview.net/forum?id=oapKSVM2bcj}
  281. }
  282. ```
  283. ## Supported python versions
  284. `einops` works with python 3.9 or later.