Upload 13 files

einops-0.4.1.dist-info/INSTALLER

@@ -0,0 +1 @@
+pip

einops-0.4.1.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Alex Rogozhnikov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

einops-0.4.1.dist-info/METADATA

@@ -0,0 +1,309 @@
+Metadata-Version: 2.1
+Name: einops
+Version: 0.4.1
+Summary: A new flavour of deep learning operations
+Home-page: https://github.com/arogozhnikov/einops
+Author: Alex Rogozhnikov
+License: UNKNOWN
+Keywords: deep learning,neural networks,tensor manipulation,machine learning,scientific computations,einops
+Platform: UNKNOWN
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3 
+Description-Content-Type: text/markdown
+License-File: LICENSE
+
+<!-- this link magically rendered as video, unfortunately not in docs -->
+
+<!-- https://user-images.githubusercontent.com/6318811/116849688-0ca41c00-aba4-11eb-8ccf-74744f6cbc23.mp4 -->
+
+<a href='http://arogozhnikov.github.io/images/einops/einops_video.mp4' >
+<div align="center">
+  <img src="http://arogozhnikov.github.io/images/einops/einops_video.gif" alt="einops package examples" />
+  <br>
+  <small><a href='http://arogozhnikov.github.io/images/einops/einops_video.mp4'>This video in high quality (mp4)</a></small>
+  <br><br>
+</div>
+</a>
+
+# einops 
+[![Run tests](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml/badge.svg)](https://github.com/arogozhnikov/einops/actions/workflows/run_tests.yml)
+[![PyPI version](https://badge.fury.io/py/einops.svg)](https://badge.fury.io/py/einops)
+[![Documentation](https://img.shields.io/badge/documentation-link-blue.svg)](https://einops.rocks/)
+![Supported python versions](https://raw.githubusercontent.com/arogozhnikov/einops/master/docs/resources/python_badge.svg)
+
+
+Flexible and powerful tensor operations for readable and reliable code. 
+Supports numpy, pytorch, tensorflow, jax, and [others](#supported-frameworks).
+
+## Recent updates:
+
+- torch.jit.script is supported for pytorch layers
+- powerful EinMix added to einops. [Einmix tutorial notebook](https://github.com/arogozhnikov/einops/blob/master/docs/3-einmix-layer.ipynb) 
+
+<!--<div align="center">
+  <img src="http://arogozhnikov.github.io/images/einops/einops_logo_350x350.png" 
+  alt="einops package logo" width="250" height="250" />
+  <br><br>
+</div> -->
+
+## Tweets 
+
+> In case you need convincing arguments for setting aside time to learn about einsum and einops...
+[Tim Rocktäschel, FAIR](https://twitter.com/_rockt/status/1230818967205425152)
+
+> Writing better code with PyTorch and einops 👌
+[Andrej Karpathy, AI at Tesla](https://twitter.com/karpathy/status/1290826075916779520)
+
+> 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
+[Nasim Rahaman, MILA (Montreal)](https://twitter.com/nasim_rahaman/status/1216022614755463169)
+
+[More testimonials](https://einops.rocks/pages/testimonials/)
+
+## Contents
+
+- [Installation](#Installation)
+- [Documentation](https://einops.rocks/)
+- [Tutorial](#Tutorials) 
+- [API micro-reference](#API)
+- [Why using einops](#Why-using-einops-notation)
+- [Supported frameworks](#Supported-frameworks)
+- [Contributing](#Contributing)
+- [Repository](https://github.com/arogozhnikov/einops) and [discussions](https://github.com/arogozhnikov/einops/discussions)
+
+## Installation  <a name="Installation"></a>
+
+Plain and simple:
+```bash
+pip install einops
+```
+
+<!--
+`einops` has no mandatory dependencies (code examples also require jupyter, pillow + backends). 
+To obtain the latest github version 
+
+```bash
+pip install https://github.com/arogozhnikov/einops/archive/master.zip
+```
+-->
+
+## Tutorials <a name="Tutorials"></a>
+
+Tutorials are the most convenient way to see `einops` in action
+
+- part 1: [einops fundamentals](https://github.com/arogozhnikov/einops/blob/master/docs/1-einops-basics.ipynb) 
+- part 2: [einops for deep learning](https://github.com/arogozhnikov/einops/blob/master/docs/2-einops-for-deep-learning.ipynb)
+- part 3: [improve pytorch code with einops](https://arogozhnikov.github.io/einops/pytorch-examples.html)   
+
+
+## API <a name="API"></a>
+
+`einops` has a minimalistic yet powerful API.
+
+Three operations provided ([einops tutorial](https://github.com/arogozhnikov/einops/blob/master/docs/) 
+shows those cover stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate, view and numerous reductions)
+
+```python
+from einops import rearrange, reduce, repeat
+# rearrange elements according to the pattern
+output_tensor = rearrange(input_tensor, 't b c -> b c t')
+# combine rearrangement and reduction
+output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2)
+# copy along a new axis 
+output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)
+```
+And two corresponding layers (`einops` keeps a separate version for each framework) with the same API.
+
+```python
+from einops.layers.chainer import Rearrange, Reduce
+from einops.layers.gluon import Rearrange, Reduce
+from einops.layers.keras import Rearrange, Reduce
+from einops.layers.torch import Rearrange, Reduce
+from einops.layers.tensorflow import Rearrange, Reduce
+```
+
+Layers behave similarly to operations and have the same parameters 
+(with the exception of the first argument, which is passed during call)
+
+```python
+layer = Rearrange(pattern, **axes_lengths)
+layer = Reduce(pattern, reduction, **axes_lengths)
+
+# apply created layer to a tensor / variable
+x = layer(x)
+```
+
+Example of using layers within a model:
+```python
+# example given for pytorch, but code in other frameworks is almost identical  
+from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
+from einops.layers.torch import Rearrange
+
+model = Sequential(
+    Conv2d(3, 6, kernel_size=5),
+    MaxPool2d(kernel_size=2),
+    Conv2d(6, 16, kernel_size=5),
+    MaxPool2d(kernel_size=2),
+    # flattening
+    Rearrange('b c h w -> b (c h w)'),  
+    Linear(16*5*5, 120), 
+    ReLU(),
+    Linear(120, 10), 
+)
+```
+
+<!---
+Additionally two auxiliary functions provided
+```python
+from einops import asnumpy, parse_shape
+# einops.asnumpy converts tensors of imperative frameworks to numpy
+numpy_tensor = asnumpy(input_tensor)
+# einops.parse_shape gives a shape of axes of interest 
+parse_shape(input_tensor, 'batch _ h w') # e.g {'batch': 64, 'h': 128, 'w': 160}
+```
+-->
+
+## Naming <a name="Naming"></a>
+
+`einops` stands for Einstein-Inspired Notation for operations 
+(though "Einstein operations" is more attractive and easier to remember).
+
+Notation was loosely inspired by Einstein summation (in particular by `numpy.einsum` operation).
+
+## Why use `einops` notation?! <a name="Why-using-einops-notation"></a>
+
+
+### Semantic information (being verbose in expectations)
+
+```python
+y = x.view(x.shape[0], -1)
+y = rearrange(x, 'b c h w -> b (c h w)')
+```
+While these two lines are doing the same job in *some* context,
+the second one provides information about the input and output.
+In other words, `einops` focuses on interface: *what is the input and output*, not *how* the output is computed.
+
+The next operation looks similar:
+
+```python
+y = rearrange(x, 'time c h w -> time (c h w)')
+```
+but it gives the reader a hint: 
+this is not an independent batch of images we are processing, 
+but rather a sequence (video). 
+
+Semantic information makes the code easier to read and maintain. 
+
+### Convenient checks
+
+Reconsider the same example:
+
+```python
+y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
+y = rearrange(x, 'b c h w -> b (c h w)')
+```
+The second line checks that the input has four dimensions, 
+but you can also specify particular dimensions. 
+That's opposed to just writing comments about shapes since 
+[comments don't work and don't prevent mistakes](https://medium.freecodecamp.org/code-comments-the-good-the-bad-and-the-ugly-be9cc65fbf83)
+as we know   
+```python
+y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
+y = rearrange(x, 'b c h w -> b (c h w)', c=256, h=19, w=19)
+```
+
+### Result is strictly determined
+
+Below we have at least two ways to define the depth-to-space operation
+```python
+# depth-to-space
+rearrange(x, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=2, w2=2)
+rearrange(x, 'b c (h h2) (w w2) -> b (h2 w2 c) h w', h2=2, w2=2)
+```
+There are at least four more ways to do it. Which one is used by the framework?
+
+These details are ignored, since *usually* it makes no difference, 
+but it can make a big difference (e.g. if you use grouped convolutions in the next stage), 
+and you'd like to specify this in your code.
+
+
+### Uniformity
+
+```python
+reduce(x, 'b c (x dx) -> b c x', 'max', dx=2)
+reduce(x, 'b c (x dx) (y dy) -> b c x y', 'max', dx=2, dy=3)
+reduce(x, 'b c (x dx) (y dy) (z dz) -> b c x y z', 'max', dx=2, dy=3, dz=4)
+```
+These examples demonstrated that we don't use separate operations for 1d/2d/3d pooling, 
+those are all defined in a uniform way. 
+
+Space-to-depth and depth-to space are defined in many frameworks but how about width-to-height? Here you go:
+
+```python
+rearrange(x, 'b c h (w w2) -> b c (h w2) w', w2=2)
+```
+
+### Framework independent behavior
+
+Even simple functions are defined differently by different frameworks
+
+```python
+y = x.flatten() # or flatten(x)
+```
+
+Suppose `x`'s shape was `(3, 4, 5)`, then `y` has shape ...
+
+- numpy, cupy, chainer, pytorch: `(60,)`
+- keras, tensorflow.layers, mxnet and gluon: `(3, 20)`
+
+`einops` works the same way in all frameworks.
+
+### Independence of framework terminology
+
+Example: `tile` vs `repeat` causes lots of confusion. To copy image along width:
+```python
+np.tile(image, (1, 2))    # in numpy
+image.repeat(1, 2)        # pytorch's repeat ~ numpy's tile
+```
+
+With einops you don't need to decipher which axis was repeated:
+```python
+repeat(image, 'h w -> h (tile w)', tile=2)  # in numpy
+repeat(image, 'h w -> h (tile w)', tile=2)  # in pytorch
+repeat(image, 'h w -> h (tile w)', tile=2)  # in tf
+repeat(image, 'h w -> h (tile w)', tile=2)  # in jax
+repeat(image, 'h w -> h (tile w)', tile=2)  # in mxnet
+... (etc.)
+```
+
+Testimonials provide user's perspective on the same question. 
+
+## Supported frameworks <a name="Supported-frameworks"></a>
+
+Einops works with ...
+
+- [numpy](http://www.numpy.org/)
+- [pytorch](https://pytorch.org/)
+- [tensorflow](https://www.tensorflow.org/)
+- [jax](https://github.com/google/jax)
+- [cupy](https://cupy.chainer.org/)
+- [chainer](https://chainer.org/)
+- [gluon](https://gluon.mxnet.io/)
+- [tf.keras](https://www.tensorflow.org/guide/keras)
+- [mxnet](https://mxnet.apache.org/) (experimental)
+
+
+## Contributing <a name="Contributing"></a>
+
+Best ways to contribute are
+
+- spread the word about `einops`
+- if you like explaining things, more tutorials/tear-downs of implementations is welcome
+- tutorials in other languages are very welcome
+- do you have project/code example to share? Let me know in github discussions
+- use `einops` in your papers!
+
+## Supported python versions
+
+`einops` works with python 3.6 or later. 
+
+

einops-0.4.1.dist-info/RECORD

@@ -0,0 +1,31 @@
+einops-0.4.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+einops-0.4.1.dist-info/LICENSE,sha256=MNmENkKW9R_67K1LAe4SfpUlDFBokY1LZvyWIGcj5DQ,1073
+einops-0.4.1.dist-info/METADATA,sha256=UdOBa4tijnwPJI48dGASJt4-czHTJ4LLiY4dfdRXffI,10737
+einops-0.4.1.dist-info/RECORD,,
+einops-0.4.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+einops-0.4.1.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+einops-0.4.1.dist-info/top_level.txt,sha256=zh9ckJ4QUP-fUBSO5-UKAcNKvC_lzMGMYS6_nnoT4Tc,7
+einops/__init__.py,sha256=8uWtV9MPDvreSsu6VG0E-j2TlxrckMGW2Zy5fCRFu6I,297
+einops/__pycache__/__init__.cpython-310.pyc,,
+einops/__pycache__/_backends.cpython-310.pyc,,
+einops/__pycache__/_torch_specific.cpython-310.pyc,,
+einops/__pycache__/einops.cpython-310.pyc,,
+einops/__pycache__/parsing.cpython-310.pyc,,
+einops/_backends.py,sha256=xyh2XysbubzGMqyRZPu8ld5-1bob693ESp6vkZR4gV8,17132
+einops/_torch_specific.py,sha256=-VuZFXozi6GPtKQFWvb7BhzswyOfjByXGG0GLvvhDXg,2804
+einops/einops.py,sha256=LkY-JdbOUqf_GtuiskedET4aaSs0yJKfVaLoBwS6UP8,27879
+einops/layers/__init__.py,sha256=JHwHQUP5sBIYhSwRrjhZYxGIdw8-UTEWUPbeEduBuBY,2824
+einops/layers/__pycache__/__init__.cpython-310.pyc,,
+einops/layers/__pycache__/_einmix.cpython-310.pyc,,
+einops/layers/__pycache__/chainer.cpython-310.pyc,,
+einops/layers/__pycache__/gluon.cpython-310.pyc,,
+einops/layers/__pycache__/keras.cpython-310.pyc,,
+einops/layers/__pycache__/tensorflow.cpython-310.pyc,,
+einops/layers/__pycache__/torch.cpython-310.pyc,,
+einops/layers/_einmix.py,sha256=k1Wt5z7KmJF9nj345ZhUXeRBcV1D2bNkz35yF82zB_E,8249
+einops/layers/chainer.py,sha256=VisqqyZiEpDl7NdCSjVSa4u7aXgZuNpA0hglkfGydiM,1927
+einops/layers/gluon.py,sha256=Ll85s1OWKqRAhSwFS33jQwbTicD1MnhrH4lbnlqvoPU,2101
+einops/layers/keras.py,sha256=RTsR-aim1Sco5VXI2W1Qs639hJRJ0hWIilTZCs3Ftn4,212
+einops/layers/tensorflow.py,sha256=xNsVaKIMoB2kZeSeFUKXq29LWz-Fppt2K2aRln5s0-Y,3269
+einops/layers/torch.py,sha256=IOdwPR2uL_ZFuzWthGz6p-8af1zg801UmjB8uTBA5HY,2379
+einops/parsing.py,sha256=75hvgp6iWvvLUe67IaQujmox1tjvF9ZsBMaXQYnQmqU,6637

einops-0.4.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.37.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

einops-0.4.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+einops

einops/__init__.py

@@ -0,0 +1,3 @@
+from .rearrange import rearrange
+
+__all__ = ["rearrange"]

einops/_parsing.py

@@ -0,0 +1,302 @@
+"""Adapted from https://github.com/arogozhnikov/einops/blob/36c7bb16e57d6e57f8f3050f9e07abdf3f00469f/einops/parsing.py.
+
+MIT License
+
+Copyright (c) 2018 Alex Rogozhnikov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+from __future__ import annotations
+
+import keyword
+import warnings
+from typing import Collection, List, Mapping, Optional, Set, Tuple, Union
+
+_ellipsis: str = "…"  # NB, this is a single unicode symbol. String is used as it is not a list, but can be iterated
+
+
+class AnonymousAxis:
+    """Used by `ParsedExpression` to represent an axis with a size (> 1), but no associated identifier.
+
+    Note: Different instances of this class are not equal to each other, even if they have the same value.
+    """
+
+    def __init__(self, value: str) -> None:
+        self.value = int(value)
+        if self.value < 1:
+            raise ValueError(
+                f"Anonymous axis should have positive length, not {self.value}"
+            )
+
+    def __repr__(self) -> str:
+        return f"{self.value}-axis"
+
+
+class ParsedExpression:
+    """Structure containing information about one side of an `einops`-style pattern (e.g. 'b c (h w)')."""
+
+    def __init__(
+        self,
+        expression: str,
+        *,
+        allow_underscore: bool = False,
+        allow_duplicates: bool = False,
+    ) -> None:
+        """Parse the expression and store relevant metadata.
+
+        Args:
+            expression (str): the `einops`-pattern to parse
+            allow_underscore (bool): whether to allow axis identifier names to begin with an underscore
+            allow_duplicates (bool): whether to allow an identifier to appear more than once in the expression
+        """
+        self.has_ellipsis: bool = False
+        self.has_ellipsis_parenthesized: Optional[bool] = None
+        self.identifiers: Set[Union[str, AnonymousAxis]] = set()
+        # that's axes like 2, 3, 4 or 5. Axes with size 1 are exceptional and replaced with empty composition
+        self.has_non_unitary_anonymous_axes: bool = False
+        # composition keeps structure of composite axes, see how different corner cases are handled in tests
+        self.composition: List[Union[List[Union[str, AnonymousAxis]], str]] = []
+        if "." in expression:
+            if "..." not in expression:
+                raise ValueError(
+                    "Expression may contain dots only inside ellipsis (...)"
+                )
+            if str.count(expression, "...") != 1 or str.count(expression, ".") != 3:
+                raise ValueError(
+                    "Expression may contain dots only inside ellipsis (...); only one ellipsis for tensor "
+                )
+            expression = expression.replace("...", _ellipsis)
+            self.has_ellipsis = True
+
+        bracket_group: Optional[List[Union[str, AnonymousAxis]]] = None
+
+        def add_axis_name(x: str) -> None:
+            if x in self.identifiers:
+                if not (allow_underscore and x == "_") and not allow_duplicates:
+                    raise ValueError(
+                        f"Indexing expression contains duplicate dimension '{x}'"
+                    )
+            if x == _ellipsis:
+                self.identifiers.add(_ellipsis)
+                if bracket_group is None:
+                    self.composition.append(_ellipsis)
+                    self.has_ellipsis_parenthesized = False
+                else:
+                    bracket_group.append(_ellipsis)
+                    self.has_ellipsis_parenthesized = True
+            else:
+                is_number = str.isdecimal(x)
+                if is_number and int(x) == 1:
+                    # handling the case of anonymous axis of length 1
+                    if bracket_group is None:
+                        self.composition.append([])
+                    else:
+                        pass  # no need to think about 1s inside parenthesis
+                    return
+                is_axis_name, reason = self.check_axis_name_return_reason(
+                    x, allow_underscore=allow_underscore
+                )
+                if not (is_number or is_axis_name):
+                    raise ValueError(f"Invalid axis identifier: {x}\n{reason}")
+                axis_name: Union[str, AnonymousAxis] = (
+                    AnonymousAxis(x) if is_number else x
+                )
+                self.identifiers.add(axis_name)
+                if is_number:
+                    self.has_non_unitary_anonymous_axes = True
+                if bracket_group is None:
+                    self.composition.append([axis_name])
+                else:
+                    bracket_group.append(axis_name)
+
+        current_identifier = None
+        for char in expression:
+            if char in "() ":
+                if current_identifier is not None:
+                    add_axis_name(current_identifier)
+                current_identifier = None
+                if char == "(":
+                    if bracket_group is not None:
+                        raise ValueError(
+                            "Axis composition is one-level (brackets inside brackets not allowed)"
+                        )
+                    bracket_group = []
+                elif char == ")":
+                    if bracket_group is None:
+                        raise ValueError("Brackets are not balanced")
+                    self.composition.append(bracket_group)
+                    bracket_group = None
+            elif str.isalnum(char) or char in ["_", _ellipsis]:
+                if current_identifier is None:
+                    current_identifier = char
+                else:
+                    current_identifier += char
+            else:
+                raise ValueError(f"Unknown character '{char}'")
+
+        if bracket_group is not None:
+            raise ValueError(f"Imbalanced parentheses in expression: '{expression}'")
+        if current_identifier is not None:
+            add_axis_name(current_identifier)
+
+    @staticmethod
+    def check_axis_name_return_reason(
+        name: str, allow_underscore: bool = False
+    ) -> Tuple[bool, str]:
+        """Check if the given axis name is valid, and a message explaining why if not.
+
+        Valid axes names are python identifiers except keywords, and should not start or end with an underscore.
+
+        Args:
+            name (str): the axis name to check
+            allow_underscore (bool): whether axis names are allowed to start with an underscore
+
+        Returns:
+            Tuple[bool, str]: whether the axis name is valid, a message explaining why if not
+        """
+        if not str.isidentifier(name):
+            return False, "not a valid python identifier"
+        elif name[0] == "_" or name[-1] == "_":
+            if name == "_" and allow_underscore:
+                return True, ""
+            return False, "axis name should should not start or end with underscore"
+        else:
+            if keyword.iskeyword(name):
+                warnings.warn(
+                    f"It is discouraged to use axes names that are keywords: {name}",
+                    RuntimeWarning,
+                )
+            if name in ["axis"]:
+                warnings.warn(
+                    "It is discouraged to use 'axis' as an axis name and will raise an error in future",
+                    FutureWarning,
+                )
+            return True, ""
+
+    @staticmethod
+    def check_axis_name(name: str) -> bool:
+        """Check if the name is a valid axis name.
+
+        Args:
+            name (str): the axis name to check
+
+        Returns:
+            bool: whether the axis name is valid
+        """
+        is_valid, _ = ParsedExpression.check_axis_name_return_reason(name)
+        return is_valid
+
+
+def parse_pattern(
+    pattern: str, axes_lengths: Mapping[str, int]
+) -> Tuple[ParsedExpression, ParsedExpression]:
+    """Parse an `einops`-style pattern into a left-hand side and right-hand side `ParsedExpression` object.
+
+    Args:
+        pattern (str): the `einops`-style rearrangement pattern
+        axes_lengths (Mapping[str, int]): any additional length specifications for dimensions
+
+    Returns:
+       Tuple[ParsedExpression, ParsedExpression]: a tuple containing the left-hand side and right-hand side expressions
+    """
+    # adapted from einops.einops._prepare_transformation_recipe
+    # https://github.com/arogozhnikov/einops/blob/230ac1526c1f42c9e1f7373912c7f8047496df11/einops/einops.py
+    try:
+        left_str, right_str = pattern.split("->")
+    except ValueError:
+        raise ValueError("Pattern must contain a single '->' separator") from None
+
+    if _ellipsis in axes_lengths:
+        raise ValueError(f"'{_ellipsis}' is not an allowed axis identifier")
+
+    left = ParsedExpression(left_str)
+    right = ParsedExpression(right_str)
+
+    if not left.has_ellipsis and right.has_ellipsis:
+        raise ValueError(
+            f"Ellipsis found in right side, but not left side of a pattern {pattern}"
+        )
+    if left.has_ellipsis and left.has_ellipsis_parenthesized:
+        raise ValueError(
+            f"Ellipsis is parenthesis in the left side is not allowed: {pattern}"
+        )
+
+    return left, right
+
+
+def validate_rearrange_expressions(
+    left: ParsedExpression, right: ParsedExpression, axes_lengths: Mapping[str, int]
+) -> None:
+    """Perform expression validations that are specific to the `rearrange` operation.
+
+    Args:
+        left (ParsedExpression): left-hand side expression
+        right (ParsedExpression): right-hand side expression
+        axes_lengths (Mapping[str, int]): any additional length specifications for dimensions
+    """
+    for length in axes_lengths.values():
+        if (length_type := type(length)) is not int:
+            raise TypeError(
+                f"rearrange axis lengths must be integers, got: {length_type}"
+            )
+
+    if left.has_non_unitary_anonymous_axes or right.has_non_unitary_anonymous_axes:
+        raise ValueError("rearrange only supports unnamed axes of size 1")
+
+    difference = set.symmetric_difference(left.identifiers, right.identifiers)
+    if len(difference) > 0:
+        raise ValueError(
+            f"Identifiers only on one side of rearrange expression (should be on both): {difference}"
+        )
+
+    unmatched_axes = axes_lengths.keys() - left.identifiers
+    if len(unmatched_axes) > 0:
+        raise ValueError(
+            f"Identifiers not found in rearrange expression: {unmatched_axes}"
+        )
+
+
+def comma_separate(collection: Collection[Union[str, Collection[str]]]) -> str:
+    """Convert a collection of strings representing first class dims into a comma-separated string.
+
+    Args:
+        collection (Collection[Union[str, Collection[str]]]): the collection of strings to convert
+
+    Returns:
+        str: the comma-separated string
+
+    Examples:
+        >>> comma_separate(('d0',))
+        'd0'
+
+        >>> comma_separate(('d0', 'd1', 'd2', 'd3'))
+        'd0, d1, d2, d3'
+
+        >>> comma_separate([('d1', 'd4')])
+        '(d1, d4)'
+
+        >>> comma_separate([('d0',), (), ('d1',), ('d2',), ('d3', 'd4')])
+        '(d0,), (), (d1,), (d2,), (d3, d4)'
+    """
+    return ", ".join(
+        item
+        if isinstance(item, str)
+        else f"({comma_separate(item)}{',' if len(item) == 1 else ''})"
+        for item in collection
+    )

einops/rearrange.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import functools
+from typing import Callable, Dict, List, Sequence, Tuple, Union
+
+import torch
+
+from functorch._C import dim as _C
+from ._parsing import (
+    _ellipsis,
+    AnonymousAxis,
+    comma_separate,
+    parse_pattern,
+    validate_rearrange_expressions,
+)
+
+__all__ = ["rearrange"]
+
+dims = _C.dims
+
+
+@functools.lru_cache(256)
+def _create_rearrange_callable(
+    tensor_ndim: int, pattern: str, **axes_lengths: int
+) -> Callable[[torch.Tensor], torch.Tensor]:
+    r"""Translate an `einops`-style pattern into a callable that performs the rearrange using first-class dimensions.
+
+    Since the an equivalent result is computed for tensors with the same number of dimensions, with the same pattern and
+    specified axes lengths, this function can be memoized.
+
+    Args:
+        tensor_ndim (int): the number of dimensions in the tensor to rearrange
+        pattern (str): the `einops`-style rearrangement pattern
+        axes_lengths (int): any additional length specifications for dimensions
+
+    Returns:
+        Callable[[torch.Tensor], torch.Tensor]: a callable that performs the rearrangement
+    """
+    left, right = parse_pattern(pattern, axes_lengths)
+    validate_rearrange_expressions(left, right, axes_lengths)
+
+    n_anon_dims = sum(not dim for dim in left.composition)
+    if left.has_ellipsis:
+        n_ellipsis_dims = tensor_ndim - (len(left.composition) - 1)
+        n_named_dims = len(left.identifiers) - 1
+
+        if (pattern_ndim := n_anon_dims + n_named_dims) > tensor_ndim:
+            raise ValueError(
+                f"Number of dimensions in pattern ({pattern_ndim}) must be less than or equal to the number of "
+                f"dimensions in the tensor ({tensor_ndim})"
+            )
+    else:
+        n_ellipsis_dims = 0
+        n_named_dims = len(left.identifiers)
+
+        if (pattern_ndim := len(left.composition)) != tensor_ndim:
+            raise ValueError(
+                f"Number of dimensions in pattern ({pattern_ndim}) must be equal to the number of dimensions in "
+                f"the tensor ({tensor_ndim})"
+            )
+    n_dims = n_named_dims + n_ellipsis_dims + n_anon_dims
+
+    if n_dims == 0:
+        # an identity rearrangement on a 0-dimension tensor
+        return lambda tensor: tensor
+
+    first_class_dims: Tuple[str, ...] = tuple(f"d{i}" for i in range(n_dims))
+    identifier_dim_map: Dict[Union[str, AnonymousAxis], Tuple[str, ...]] = {}
+    anon_axes: List[AnonymousAxis] = []
+
+    # map the left-hand side identifiers to strings representing first class dims
+    dims_i = 0
+    for dimension in left.composition:
+        if isinstance(dimension, list):
+            for identifier in dimension:
+                # non-unitary anon axes are not allowed in rearrange & unitary anon axes are represented as empty lists
+                assert isinstance(identifier, str)
+                identifier_dim_map[identifier] = (first_class_dims[dims_i],)
+                dims_i += 1
+            if not dimension:
+                # unitary anonymous axis
+                anon_axis = AnonymousAxis("1")
+                identifier_dim_map[anon_axis] = (first_class_dims[dims_i],)
+                anon_axes.append(anon_axis)
+                dimension.append(anon_axis)
+                dims_i += 1
+        elif dimension == _ellipsis:
+            identifier = _ellipsis
+            identifier_dim_map[identifier] = tuple(
+                first_class_dims[dims_i + j] for j in range(n_ellipsis_dims)
+            )
+            dims_i += n_ellipsis_dims
+        else:
+            raise ValueError(f"Unexpected dimension: {dimension}")
+
+    def composition_to_dims(
+        composition: Sequence[Union[List[Union[str, AnonymousAxis]], str]]
+    ) -> List[Union[str, Tuple[str, ...]]]:
+        """Convert a `ParsedExpression.composition` into a `Tensor.__getitem__` index of strings representing first
+        class dims."""
+        dim_composition: List[Union[str, Tuple[str, ...]]] = []
+        for dimension in composition:
+            if isinstance(dimension, list):
+                dim_composition.append(
+                    tuple(
+                        dim
+                        for identifier in dimension
+                        for dim in identifier_dim_map[identifier]
+                    )
+                )
+            elif dimension == _ellipsis:
+                dim_composition.extend(identifier_dim_map[_ellipsis])
+            else:
+                raise ValueError(f"Unexpected dimension: {dimension}")
+        return dim_composition
+
+    left_dims = composition_to_dims(left.composition)
+    right_dims = composition_to_dims(right.composition)
+    anon_dims = tuple(identifier_dim_map[axis][0] for axis in anon_axes)
+    specified_lengths = tuple(
+        (identifier_dim_map[axis][0], length) for axis, length in axes_lengths.items()
+    )
+
+    custom_rearrange_callable_name = "do_rearrange"
+    custom_rearrange_callable_code = (
+        (
+            f"def {custom_rearrange_callable_name}(tensor):\n"
+            f"    {comma_separate(first_class_dims)} = dims({n_dims})\n"
+        )
+        + (
+            "".join(
+                f"    {dim}.size = {length}\n" for (dim, length) in specified_lengths
+            )
+            if specified_lengths
+            else ""
+        )
+        + f"    tensor = tensor[{comma_separate(left_dims)}].order({comma_separate(right_dims)})\n"
+        + (
+            f"    return tensor.sum({comma_separate([anon_dims])}, keepdim=False)\n"
+            if anon_dims
+            else "    return tensor\n"
+        )
+    )
+
+    exec(custom_rearrange_callable_code)
+    return locals()[custom_rearrange_callable_name]
+
+
+def rearrange(
+    tensor: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor, ...]],
+    pattern: str,
+    **axes_lengths: int,
+) -> torch.Tensor:
+    r"""A native implementation of `einops.rearrange`, a reader-friendly smart element reordering for multidimensional
+    tensors. This operation includes functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze,
+    stack, concatenate and other operations.
+
+    See: https://einops.rocks/api/rearrange/
+
+    Args:
+        tensor (Tensor or sequence of Tensor): the tensor(s) to rearrange
+        pattern (str): the rearrangement pattern
+        axes_lengths (int): any additional length specifications for dimensions
+
+    Returns:
+        Tensor: the rearranged tensor
+
+    Examples:
+        >>> # suppose we have a set of 32 images in "h w c" format (height-width-channel)
+        >>> images = torch.randn((32, 30, 40, 3))
+
+        >>> # stack along first (batch) axis, output is a single array
+        >>> rearrange(images, 'b h w c -> b h w c').shape
+        torch.Size([32, 30, 40, 3])
+
+        >>> # concatenate images along height (vertical axis), 960 = 32 * 30
+        >>> rearrange(images, 'b h w c -> (b h) w c').shape
+        torch.Size([960, 40, 3])
+
+        >>> # concatenated images along horizontal axis, 1280 = 32 * 40
+        >>> rearrange(images, 'b h w c -> h (b w) c').shape
+        torch.Size([30, 1280, 3])
+
+        >>> # reordered axes to "b c h w" format for deep learning
+        >>> rearrange(images, 'b h w c -> b c h w').shape
+        torch.Size([32, 3, 30, 40])
+
+        >>> # flattened each image into a vector, 3600 = 30 * 40 * 3
+        >>> rearrange(images, 'b h w c -> b (c h w)').shape
+        torch.Size([32, 3600])
+
+        >>> # split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2
+        >>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
+        torch.Size([128, 15, 20, 3])
+
+        >>> # space-to-depth operation
+        >>> rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape
+        torch.Size([32, 15, 20, 12])
+    """
+    if not isinstance(tensor, torch.Tensor):
+        tensor = torch.stack(tensor)
+
+    rearrange_callable = _create_rearrange_callable(
+        tensor.ndim, pattern, **axes_lengths
+    )
+
+    return rearrange_callable(tensor)