{"id":13411834,"url":"https://github.com/arogozhnikov/einops","last_synced_at":"2025-05-12T13:04:17.601Z","repository":{"id":37336057,"uuid":"149832847","full_name":"arogozhnikov/einops","owner":"arogozhnikov","description":"Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)","archived":false,"fork":false,"pushed_at":"2025-04-25T09:16:18.000Z","size":2726,"stargazers_count":8882,"open_issues_count":32,"forks_count":369,"subscribers_count":65,"default_branch":"main","last_synced_at":"2025-05-03T10:52:42.498Z","etag":null,"topics":["chainer","cupy","deep-learning","einops","jax","keras","numpy","pytorch","tensor","tensorflow"],"latest_commit_sha":null,"homepage":"https://einops.rocks","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/arogozhnikov.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2018-09-22T00:45:08.000Z","updated_at":"2025-05-02T02:29:37.000Z","dependencies_parsed_at":"2024-01-11T10:29:04.699Z","dependency_job_id":"c3e1d165-999f-44eb-a200-d5b920ff719d","html_url":"https://github.com/arogozhnikov/einops","commit_stats":{"total_commits":616,"total_committers":29,"mean_commits":21.24137931034483,"dds":"0.19967532467532467","last_synced_commit":"1d3f9c48047ba147211c8f15a079db8e87766a0b"},"previous_names":[],"tags_count":16,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arogozhnikov%2Feinops","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arogozhnikov%2Feinops/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arogozhnikov%2Feinops/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arogozhnikov%2Feinops/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/arogozhnikov","download_url":"https://codeload.github.com/arogozhnikov/einops/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":252332841,"owners_count":21731086,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["chainer","cupy","deep-learning","einops","jax","keras","numpy","pytorch","tensor","tensorflow"],"created_at":"2024-07-30T20:01:17.381Z","updated_at":"2025-05-04T12:27:23.742Z","avatar_url":"https://github.com/arogozhnikov.png","language":"Python","funding_links":[],"categories":["Python","机器学习框架","Linear Algebra / Statistics Toolkit","General","Computation and Communication Optimisation","PyTorch Tools, Libraries, and Frameworks","👁 Vision \u0026 ViT Ecosystem","Other","Frameworks","Deep Learning Tools","1. Core Frameworks \u0026 Libraries"],"sub_categories":["General Purpose Tensor Library","General-Purpose Machine Learning"],"readme":"\n\u003c!--\n\u003ca href='http://arogozhnikov.github.io/images/einops/einops_video.mp4' \u003e\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"http://arogozhnikov.github.io/images/einops/einops_video.gif\" alt=\"einops package examples\" /\u003e\n  \u003cbr\u003e\n  \u003csmall\u003e\u003ca href='http://arogozhnikov.github.io/images/einops/einops_video.mp4'\u003eThis video in high quality (mp4)\u003c/a\u003e\u003c/small\u003e\n  \u003cbr\u003e\u003cbr\u003e\n\u003c/div\u003e\n\u003c/a\u003e\n--\u003e\n\n\u003c!-- this link magically rendered as video on github readme, unfortunately not in docs --\u003e\n\nhttps://user-images.githubusercontent.com/6318811/177030658-66f0eb5d-e136-44d8-99c9-86ae298ead5b.mp4\n\n\n\n\n# einops \n[![Run tests](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml/badge.svg)](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml)\n[![PyPI version](https://badge.fury.io/py/einops.svg)](https://badge.fury.io/py/einops)\n[![Documentation](https://img.shields.io/badge/documentation-link-blue.svg)](https://einops.rocks/)\n![Supported python versions](https://raw.githubusercontent.com/arogozhnikov/einops/main/docs/resources/python_badge.svg)\n\n\nFlexible and powerful tensor operations for readable and reliable code. \u003cbr /\u003e\nSupports numpy, pytorch, tensorflow, jax, and [others](#supported-frameworks).\n\n## Recent updates:\n\n- 0.8.0: tinygrad backend added, small fixes\n- 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\n- 10'000🎉: github reports that more than 10k project use einops\n- einops 0.6.1: paddle backend added\n- einops 0.6 introduces [packing and unpacking](https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb)\n- einops 0.5: einsum is now a part of einops\n- [Einops paper](https://openreview.net/pdf?id=oapKSVM2bcj) is accepted for oral presentation at ICLR 2022 (yes, it worth reading).\n  Talk recordings are [available](https://iclr.cc/virtual/2022/oral/6603)\n\n\n\u003cdetails markdown=\"1\"\u003e\n\u003csummary\u003ePrevious updates\u003c/summary\u003e\n- flax and oneflow backend added\n- torch.jit.script is supported for pytorch layers\n- powerful EinMix added to einops. [Einmix tutorial notebook](https://github.com/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb) \n\u003c/details\u003e\n\n\u003c!--\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"http://arogozhnikov.github.io/images/einops/einops_logo_350x350.png\" \n  alt=\"einops package logo\" width=\"250\" height=\"250\" /\u003e\n  \u003cbr\u003e\u003cbr\u003e\n\u003c/div\u003e --\u003e\n\n\n## Tweets \n\n\u003e In case you need convincing arguments for setting aside time to learn about einsum and einops...\n[Tim Rocktäschel](https://twitter.com/_rockt/status/1230818967205425152)\n\n\u003e Writing better code with PyTorch and einops 👌\n[Andrej Karpathy](https://twitter.com/karpathy/status/1290826075916779520)\n\n\u003e 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\n[Nasim Rahaman](https://twitter.com/nasim_rahaman/status/1216022614755463169)\n\n[More testimonials](https://einops.rocks/pages/testimonials/)\n\n\n## Contents\n\n- [Installation](#Installation)\n- [Documentation](https://einops.rocks/)\n- [Tutorial](#Tutorials)\n- [API micro-reference](#API)\n- [Why use einops](#Why-use-einops-notation)\n- [Supported frameworks](#Supported-frameworks)\n- [Citing](#Citing)\n- [Repository](https://github.com/arogozhnikov/einops) and [discussions](https://github.com/arogozhnikov/einops/discussions)\n\n## Installation  \u003ca name=\"Installation\"\u003e\u003c/a\u003e\n\nPlain and simple:\n```bash\npip install einops\n```\n\n## Tutorials \u003ca name=\"Tutorials\"\u003e\u003c/a\u003e\n\nTutorials are the most convenient way to see `einops` in action\n\n- part 1: [einops fundamentals](https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb)\n- part 2: [einops for deep learning](https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb)\n- part 3: [packing and unpacking](https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb)\n- part 4: [improve pytorch code with einops](http://einops.rocks/pytorch-examples.html)\n\nKapil Sachdeva recorded a small [intro to einops](https://www.youtube.com/watch?v=xGy75Pjsqzo).\n\n## API \u003ca name=\"API\"\u003e\u003c/a\u003e\n\n`einops` has a minimalistic yet powerful API.\n\nThree core operations provided ([einops tutorial](https://github.com/arogozhnikov/einops/blob/main/docs/)\nshows those cover stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate, view and numerous reductions)\n\n```python\nfrom einops import rearrange, reduce, repeat\n# rearrange elements according to the pattern\noutput_tensor = rearrange(input_tensor, 't b c -\u003e b c t')\n# combine rearrangement and reduction\noutput_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -\u003e b h w c', 'mean', h2=2, w2=2)\n# copy along a new axis\noutput_tensor = repeat(input_tensor, 'h w -\u003e h w c', c=3)\n```\n\nLater additions to the family are `pack` and `unpack` functions (better than stack/split/concatenate):\n\n```python\nfrom einops import pack, unpack\n# pack and unpack allow reversibly 'packing' multiple tensors into one.\n# Packed tensors may be of different dimensionality:\npacked,  ps = pack([class_token_bc, image_tokens_bhwc, text_tokens_btc], 'b * c')\nclass_emb_bc, image_emb_bhwc, text_emb_btc = unpack(transformer(packed), ps, 'b * c')\n```\n\nFinally, einops provides einsum with a support of multi-lettered names:\n\n```python\nfrom einops import einsum, pack, unpack\n# einsum is like ... einsum, generic and flexible dot-product\n# but 1) axes can be multi-lettered  2) pattern goes last 3) works with multiple frameworks\nC = einsum(A, B, 'b t1 head c, b t2 head c -\u003e b head t1 t2')\n```\n\n### EinMix\n\n`EinMix` is a generic linear layer, perfect for MLP Mixers and similar architectures.\n\n### Layers\n\nEinops provides layers (`einops` keeps a separate version for each framework) that reflect corresponding functions\n\n```python\nfrom einops.layers.torch      import Rearrange, Reduce\nfrom einops.layers.tensorflow import Rearrange, Reduce\nfrom einops.layers.flax       import Rearrange, Reduce\nfrom einops.layers.paddle     import Rearrange, Reduce\n```\n\n\u003cdetails markdown=\"1\"\u003e\n\u003csummary\u003eExample of using layers within a pytorch model\u003c/summary\u003e\nExample given for pytorch, but code in other frameworks is almost identical\n\n```python \nfrom torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU\nfrom einops.layers.torch import Rearrange\n\nmodel = Sequential(\n    ...,\n    Conv2d(6, 16, kernel_size=5),\n    MaxPool2d(kernel_size=2),\n    # flattening without need to write forward\n    Rearrange('b c h w -\u003e b (c h w)'),\n    Linear(16*5*5, 120),\n    ReLU(),\n    Linear(120, 10),\n)\n```\n\nNo more flatten needed!\n\nAdditionally, torch layers as those are script-able and compile-able.\nOperations [are torch.compile-able](https://github.com/arogozhnikov/einops/wiki/Using-torch.compile-with-einops),\n but not script-able due to limitations of torch.jit.script.\n\u003c/details\u003e\n\n\n\n\n## Naming \u003ca name=\"Naming\"\u003e\u003c/a\u003e\n\n`einops` stands for Einstein-Inspired Notation for operations \n(though \"Einstein operations\" is more attractive and easier to remember).\n\nNotation was loosely inspired by Einstein summation (in particular by `numpy.einsum` operation).\n\n## Why use `einops` notation?! \u003ca name=\"Why-use-einops-notation\"\u003e\u003c/a\u003e\n\n\n### Semantic information (being verbose in expectations)\n\n```python\ny = x.view(x.shape[0], -1)\ny = rearrange(x, 'b c h w -\u003e b (c h w)')\n```\nWhile these two lines are doing the same job in *some* context,\nthe second one provides information about the input and output.\nIn other words, `einops` focuses on interface: *what is the input and output*, not *how* the output is computed.\n\nThe next operation looks similar:\n\n```python\ny = rearrange(x, 'time c h w -\u003e time (c h w)')\n```\nbut it gives the reader a hint:\nthis is not an independent batch of images we are processing,\nbut rather a sequence (video).\n\nSemantic information makes the code easier to read and maintain.\n\n### Convenient checks\n\nReconsider the same example:\n\n```python\ny = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)\ny = rearrange(x, 'b c h w -\u003e b (c h w)')\n```\nThe second line checks that the input has four dimensions,\nbut you can also specify particular dimensions.\nThat's opposed to just writing comments about shapes since comments don't prevent mistakes,\nnot tested, and without code review tend to be outdated\n```python\ny = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)\ny = rearrange(x, 'b c h w -\u003e b (c h w)', c=256, h=19, w=19)\n```\n\n### Result is strictly determined\n\nBelow we have at least two ways to define the depth-to-space operation\n```python\n# depth-to-space\nrearrange(x, 'b c (h h2) (w w2) -\u003e b (c h2 w2) h w', h2=2, w2=2)\nrearrange(x, 'b c (h h2) (w w2) -\u003e b (h2 w2 c) h w', h2=2, w2=2)\n```\nThere are at least four more ways to do it. Which one is used by the framework?\n\nThese details are ignored, since *usually* it makes no difference,\nbut it can make a big difference (e.g. if you use grouped convolutions in the next stage),\nand you'd like to specify this in your code.\n\n\n### Uniformity\n\n```python\nreduce(x, 'b c (x dx) -\u003e b c x', 'max', dx=2)\nreduce(x, 'b c (x dx) (y dy) -\u003e b c x y', 'max', dx=2, dy=3)\nreduce(x, 'b c (x dx) (y dy) (z dz) -\u003e b c x y z', 'max', dx=2, dy=3, dz=4)\n```\nThese examples demonstrated that we don't use separate operations for 1d/2d/3d pooling,\nthose are all defined in a uniform way.\n\nSpace-to-depth and depth-to space are defined in many frameworks but how about width-to-height? Here you go:\n\n```python\nrearrange(x, 'b c h (w w2) -\u003e b c (h w2) w', w2=2)\n```\n\n\n### Framework independent behavior\n\nEven simple functions are defined differently by different frameworks\n\n```python\ny = x.flatten() # or flatten(x)\n```\n\nSuppose `x`'s shape was `(3, 4, 5)`, then `y` has shape ...\n\n- numpy, pytorch, cupy, chainer, jax: `(60,)`\n- keras, tensorflow.layers, gluon: `(3, 20)`\n\n`einops` works the same way in all frameworks.\n\n\n### Independence of framework terminology\n\nExample: `tile` vs `repeat` causes lots of confusion. To copy image along width:\n```python\nnp.tile(image, (1, 2))    # in numpy\nimage.repeat(1, 2)        # pytorch's repeat ~ numpy's tile\n```\n\nWith einops you don't need to decipher which axis was repeated:\n```python\nrepeat(image, 'h w -\u003e h (tile w)', tile=2)  # in numpy\nrepeat(image, 'h w -\u003e h (tile w)', tile=2)  # in pytorch\nrepeat(image, 'h w -\u003e h (tile w)', tile=2)  # in tf\nrepeat(image, 'h w -\u003e h (tile w)', tile=2)  # in jax\nrepeat(image, 'h w -\u003e h (tile w)', tile=2)  # in cupy\n... (etc.)\n```\n\n[Testimonials](https://einops.rocks/pages/testimonials/) provide users' perspective on the same question.\n\n\n## Supported frameworks \u003ca name=\"Supported-frameworks\"\u003e\u003c/a\u003e\n\nEinops works with ...\n\n- [numpy](http://www.numpy.org/)\n- [pytorch](https://pytorch.org/)\n- [tensorflow](https://www.tensorflow.org/)\n- [jax](https://github.com/google/jax)\n- [cupy](https://github.com/cupy/cupy)\n- [flax](https://github.com/google/flax) (community)\n- [paddle](https://github.com/PaddlePaddle/Paddle) (community)\n- [oneflow](https://github.com/Oneflow-Inc/oneflow) (community)\n- [tinygrad](https://github.com/tinygrad/tinygrad) (community)\n- [pytensor](https://github.com/pymc-devs/pytensor) (community)\n\nAdditionally, einops can be used with any framework that supports\n[Python array API standard](https://data-apis.org/array-api/latest/API_specification/index.html),\nwhich includes\n\n- numpy \u003e= 2.0\n- [MLX](https://github.com/ml-explore/mlx)  # yes, einops works with apple's framework\n- [pydata/sparse](https://github.com/pydata/sparse) \u003e= 0.15 # and works with sparse tensors\n- [cubed](https://github.com/cubed-dev/cubed) # and with distributed tensors too\n- [quantco/ndonnx](https://github.com/Quantco/ndonnx)\n- recent releases of jax and cupy.\n- dask is supported via [array-api-compat](https://github.com/data-apis/array-api-compat)\n\n\n## Development\n\nDevcontainer is provided, this environment can be used locally, or on your server,\nor within github codespaces. \nTo start with devcontainers in vs code, clone repo, and click 'Reopen in Devcontainer'. \n\nStarting from einops 0.8.1, einops distributes tests as a part of package.\n\n```bash\n# pip install einops pytest\npython -m einops.tests.run_tests numpy pytorch jax --pip-install\n```\n\n`numpy pytorch jax` is an _example_, any subset of testable frameworks can be provided.\nEvery framework is tested against numpy, so it is a requirement for tests.\n\nSpecifying `--pip-install` will install requirements in current virtualenv,\nand should be omitted if dependencies are installed locally.\n\nTo build/test docs:\n\n```bash\nhatch run docs:serve  # Serving on http://localhost:8000/\n```\n\n\n## Citing einops \u003ca name=\"Citing\"\u003e\u003c/a\u003e\n\nPlease use the following bibtex record\n\n```text\n@inproceedings{\n    rogozhnikov2022einops,\n    title={Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation},\n    author={Alex Rogozhnikov},\n    booktitle={International Conference on Learning Representations},\n    year={2022},\n    url={https://openreview.net/forum?id=oapKSVM2bcj}\n}\n```\n\n\n## Supported python versions\n\n`einops` works with python 3.9 or later.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farogozhnikov%2Feinops","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Farogozhnikov%2Feinops","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farogozhnikov%2Feinops/lists"}