Hey! Loving einops, so much that now I feel a bit sad about standard einsum not being able to use descriptive names for dimensions. It would be amazing if einops implemented einsum with the same conveniences.

This seems like it should be (intuitively) plausible:

rearrange(x, 'b -> a b c', a=1, c=1)

to essentially push a vector to be compatible with some other tensors (for broadcasting operations). Currently this throws an error.

One (sort of ugly) workaround is:

rearrange(x, '(a b c) -> a b c', a=1, c=1)

However, it seems like this is a bit redundant and it obfuscates the intent a bit. Thoughts?

Thank you for making our life easier when working with tensors. I have the following suggestions based on #50 and #20.

A. Reorder and concatenation of items of different shapes

A.1 Reorder elements of 1 dimension

As suggested in #50, it is indeed useful when we have an operation for reordering the elements of channels, especially for those working on images with different libraries (open-cv, PIL). It is really better than doing with boring indices.

I totally agree with @remisphere that we can use reorder without misleading to users.

# instead of doing this
out = imgs[:, [2, 0, 1, 3], :, : ]
# we can use the below
einops.reorder(imgs, 'batch [rg b a -> b rg a] h w', rg=2, b=1, a=1)

A.2 Concatenation of items of different sizes on 1 dimension

Since we only perform operations on the single dimension, we can perform the concatenation of multiple items with different sizes on that dimension. This will easily handle the case mentioned in #20 and extremely useful for those who use concatenate in their code. I use this function many times to concatenate tensors of different shapes. For example:

# three below tensors have different size on the 2nd dim
print(x.shape) # [b, 10]
print(y.shape) # [b, 15]
print(z.shape) # [b, 20]

# we can concatenate them as
inputs = [x, y, z]
out = einops.reorder(inputs, 'batch [x y z -> x y z]', x=10, y=15, z=20)

The above call is consistent with einops.rearrange to concatenate inputs including items of the same shape.

It is possible to split out into their components x, y, z with three lines using the below chunk function:

x = einops.chunk(out, 'batch [x yz -> x]', x=10)
y = einops.chunk(out, 'batch [x y z -> y]', x=10, y=15)
z = einops.chunk(out, 'batch [xy z -> z]', z=20)

B. Chunking along 1 dimension

In contrast with #50, I don't think it is a good idea to merge chunking into reorder. We can separate these functionalities into the above reorder and chunk. Chunking is used frequently when we want to sample parts of datasets and features.

Example in #50:

# remove the alpha channel and the bottom half of 256*256 images:
einops.chunk(imgs, 'batch [rg b a -> b rg] [top bottom -> top] w', rg=2, b=1, top=128, batch=10)

Split dataset into train and val

train_len = int(len(dataset) * 0.8)
train_split = einops.chunk(dataset, '[train val -> train] c h w', train=train_len)
val_split = einops.chunk(dataset, '[train val -> val] c h w', train=train_len)

And we can get the full dataset given train_split and val_split:

dataset = einops.reorder([train_split, val_split], '[train val -> train val] c h w', train=len(train_split), val=len(val_split))

This creates the functional einsum function as requested on #73. CC @arogozhnikov @cgarciae

The current implementation simply parses the string and converts it to einsum notation by mapping axis names to single characters (I use string.ascii_letters, starting from a, b, c etc).

Currently, it has the following features:

• Supports the backends: tensorflow, numpy, jax, pytorch, chainer, oneflow, keras, cupy.
• Allows for an arbitrary number of tensors passed.
• Allows ellipsis specification, including for multiple tensors, so long as it is provided on both the left and the right of the ->.

It does not currently support

• Reshape operations, such as "(batch channel) feature, feature -> batch channel".
• Custom reduction operations.

These could be added later if desired. Some backends do not support custom reductions in their einsum implementations so it will be a bit more work.

I also added a docstring and some unittests (in tests/test_einsum.py).

Here are some examples of use, with the numpy backend:

# Filter a set of images:
>>> batched_images = np.random.randn(128, 16, 16)
>>> filters = np.random.randn(16, 16, 30)
>>> result = einsum(batched_images, filters,
...                 "batch h w, h w channel -> batch channel") 

>>> result.shape
(128, 30)

# Matrix multiplication, with an unknown input shape:
>>> batch_shape = (50, 30)
>>> data = np.random.randn(*batch_shape, 20)
>>> weights = np.random.randn(10, 20)
>>> result = einsum(weights, data, 
...                 "out_dim in_dim, ... in_dim -> ... out_dim")
>>> result.shape
(50, 30, 10)

Note that the number of spaces next to the comma above are arbitrary, you could do either "in_dim, ..." or "in_dim , ..." - both will work.

Eager to hear feedback on this!

Cheers, Miles

Edit 1: Got working for repeat indices on one side (as used in, e.g., trace). Edit 2: Added support for chainer, oneflow, cupy, tensorflow.keras. Edit 3: Added many more tests, some mirroring those used in the np.einsum tests. Edit 4: More and more unit tests. Edit 5: Tweaked the syntax to have tensors first, pattern second. Adapted tests, and added new validation for order of arguments.

EinMix layer looks great but cannot be used with Jax since it's not a function. Would be great to have that EinMix layer in a Jax-based framework like Flax.

I love your project, it has really changed the way I write my code. I wanted to be able to use it in the TensorFlow version of Keras as well. I only had to change one thing since the dimensions are represented differently.

Describe the bug

MR #211 add py.typed to enable type checking support. However the library is not fully typed (at least not the public interface, see for instance https://github.com/arogozhnikov/einops/blob/master/einops/layers/torch.py.

Now mypy is complaining on each einops.layers.torch class call (Reduce, Rearrange, etc.).

-> error: Call to untyped function "Rearrange" in typed context [no-untyped-call]

The library should not have a py.typed file until its public interface is fully typed.

Reproduction steps run mypy on any source file containing, for instance, from einops.layers.torch import Rearrange

Expected behavior mypy should not raise [no-untyped-call] errors

Your platform Ubuntu 22.04

Updated The original issue was closed, despite the fact that provided snippet still uses a module, as opposed to a function and tested with newer and older PyTorch versions. Opening the issue here in hope of restarting the conversation and finding possible solutions.

Describe the bug The einops package does not support torch.jit.script. The following snippet tries a simple Rearrange operation, using a module to test this. It is observed that the operation is not supported.

Reproduction steps The following snippet should illustrate a concise way of reproducing this issue:

import torch
import torch.nn as nn
from einops.layers.torch import Rearrange

class SimpleRearrange(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Sequential(Rearrange('b c h w -> b h w c'))
    
    def forward(self, x):
        result = self.layer(x)
        return result
    
net = SimpleRearrange()
net.eval()
with torch.no_grad():
    torch.jit.script(net)

Expected behavior This is the expected output:

Your platform einops version: 0.3.2 Python version: 3.8.12 PyTorch version: 1.10.0 CUDA version: 10.2

Adds Flax Module's for einops operations. More details can be found in #153

einops supports JAX arrays, so I'm unsure if I should add tests or not ?

CC: @arogozhnikov

Hi Alex again! :wave: Thank you again for this wonderful work :pray: I still can't get over how beautiful the abstraction is, and how much of a multiplier it is in my work.

I was wondering what your thoughts are about this potential extension to einops? A common pattern I run into is the need to flatten dimensions (for some operation), and then to reconstitute it later. To do this, I usually need to save the dimensions to a list or tuple, and then later run a reshape on the output. To do this with einops rearrange currently is difficult, as it is unable to take in a list of dimensions (it needs each dimension to be explicitly named), and there is no way to name the ellipsis-es.

I hacked up a prototype here https://github.com/lucidrains/memorizing-transformers-pytorch/commit/b190c5ec6144d6dafb11eb356fe51dd10bccc052 to give you an idea of the use-case. Would be curious to hear your thoughts, whether you can think of some even more general way to handle this with the internal ellipsis code you have already, or whether you think this doesn't belong in einops proper

Thank you!

Describe the bug Thanks for this great work. In recent updates, it is mentioned that einops supports torch.jit.script for PyTorch layers. I and @yiheng-wang-nv have been looking into this to support TorchScript for a number of models. However, we are not able to utilize this functionality for simple operations such as Rearrange.

Reproduction steps The following snippet should illustrate a concise way of reproducing this issue:

import torch
import torch.nn as nn
from einops.layers.torch import Rearrange

class SimpleRearrange(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Sequential(Rearrange('b c h w -> b h w c'))
    
    def forward(self, x):
        result = self.layer(x)
        return result
    
net = SimpleRearrange()
net.eval()
with torch.no_grad():
    torch.jit.script(net)

Expected behavior This is the expected output:

Your platform einops version: 0.3.2 Python version: 3.8.12 PyTorch version: 1.6.0 CUDA version: 10.2

Based on this, I believe einops does not support torch.jit.script, unless we are missing something. Appreciate your inputs here.

Are there instructions for how to set up the project for contributing? The link on the README is broken. I tried building the project locally, but I had trouble with dependencies; the pyproject.toml file doesn't list any dependencies, but the tests rely on e.g. numpy. Sorry if this is a stupid question; I'm pretty new to this!

Hi,

The API design, documentation, and examples in this repository are so good that they do not need any other tutorial.

But to make others aware of the existence of this fantastic work I have created a tutorial on it. https://www.youtube.com/watch?v=xGy75Pjsqzo

I have only talked about the rearrange API and maybe later would cover other goodies from this package.

Many thanks for this great work

Regards Kapil

1. Try to collect use-cases If you don't have access to the input API to use pack but still want to unpack it with einops, you would need to know the size of the output. This is not ideal if you known the proportion that you want to unpack into.

x = MyBigPretrainedLibraryImportModel(size='base')  # base, large, huge 

# Instead of
b, s, c = x.shape  # c is different depending on model size
c_param = c // 2
x_mean, x_logvar = unpack(x, [[c_param], [c_param]], 'b s *')

# Accept this
x_mean, x_logvar = unpack(x, [[0.5], [0.5]], 'b s *')

3. Integrity - does it interplay well with existing operations and notation in einops? pack and unpack currently only accepts int, so accepting float shouldn't break anything. Incase of float it will be a bit more job, like ensuring all sum to 1, each proportion will scale up to a proper integer etc.

assert c == 512
x1, x_2, x_3 = unpack(x, [[1/3], [1/3], [1/4], 'b s *')  # Raise error as 1/3 + 1/3 + 1/4 != 1
x1, x_2, x_3 = unpack(x, [[1/3], [1/3], [1/3], 'b s *')  # Raise error as 512 is not divisible by 3

4. Readability The complexity of checking validity of proportion might not worth adding it to the library and should be handled by a few lines of user code.

Francois Fleuret suggested that it'd be nice if there were an function whereby einops.rearrange could be easily 'inverted' i.e. undone or "transformed back".

I replied that a "wrapper" function or class shouldn't be too hard, and wrote one at the following link, which includes a few examples: https://gist.github.com/drscotthawley/81865a5c5e729b769486efb9c3f2249d

Whether such a functionality just remains as an "external wrapper" that users can add-on, or somehow gets added to the einops codebase (maybe not as this class, but something similar) is up to you, but wanted to share it here to add to the conversation!

One way to include it into the existing codebase could be to rename einops.rearrange to einops._rearrange and then have the new einops.rearrange =RearrangeWrapper() where RearrangeWrapper's sub-methods call _rearrange (as shown in my gist example). If that sounds interesting then I could submit a PR.

Hello, it would be really nice if one could name the eniops layers just like any other keras layer. Right now, the following code triggers an error.

from einops.layers.tensorflow import Rearrange

tf.keras.Sequential([
    tf.keras.Input((224, 224, 3), name="inputs"),
    Rearrange("b h w c -> b c h w", name="rearrange_layer_1"),
])

The error goes away if we do not name the einops layer.

from einops.layers.tensorflow import Rearrange

tf.keras.Sequential([
    tf.keras.Input((224, 224, 3), name="inputs"),
    Rearrange("b h w c -> b c h w"),
])

Naming layers is very useful in keras, specially when using Functional models, to extract intermediate representations or to add new nodes to the graph. This process of extracting nodes is done by accessing the model's layer with model.get_layer(name).