Create einsum operation

[MilesCranmer]

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.

[MilesCranmer]

I think implementing the rearrange operations inside shouldn't be too bad: since einsum doesn't require any specific order of indices, you could call

for (tensor, left_expression) in zip(tensors, left_expressions):
    axis_names = ...
    tensor = rearrange(tensor, left_expression + "->" + " ".join(axis_names))

for every left_expression. Then, you would pass " ".join(axis_names) back to einops.einsum.

Then, you could do a similar rearrange on the output expression.

What do you think @arogozhnikov? (for a future PR, of course)

[MilesCranmer]

Actually I realized this doesn't work with repeat variable names. e.g., if I want to compute the trace of a tensor:

einsum("index index -> ", np.ones((5, 5)))

ParsedExpression doesn't allow duplicate dimensions so this doesn't work. I guess I could modify it to allow this.

Edit: fixed with allow_duplicates argument to ParsedExpression.

[arogozhnikov]

cool, thanks Miles for very clean PR and collecting test suite

I've left some thoughts and request for two kinds of tests:

1. fail on parsing of features that aren't yet supported
2. add tests for symbolic backends. Probably latter would only include tf, but still

[MilesCranmer]

Okay, all suggestions implemented. Let me know what you think.

[MilesCranmer]

Okay everything is implemented for the new syntax:

y = einsum(x, x, "i j, i k -> j k")

I also added new validation checks for the argument order, and corresponding unit-tests.

In the unittests, I also now check the specific message of each error, rather than the error type.

Let me know what you think.

[arogozhnikov]

for every left_expression. Then, you would pass " ".join(axis_names) back to einops.einsum.
Then, you could do a similar rearrange on the output expression.

It's trickier since you want some of axes to be derived from the inputs shapes.
In the examples you previously posted it could be like (i j) k, i -> j k, so first shape of the second argument should be parsed. For the output that's actually straighforward and just applying a pattern would always work

Not relevant for the PR. Just commenting since you asked

[arogozhnikov]

want to understand why it looks so strange in tf.keras

[MilesCranmer]

I'm not sure if I was interpreting the symbolic (layer=True) backends correctly or not.

Basically, this einsum assumes the x tensors have a leading batch axis, which are assumed to not be specified in the pattern. I assumed this because the create_symbol method specifies the shape as a batch shape, rather than an absolute shape. Is that correct, or should it assume the pattern also specifies the batch axis?

[MilesCranmer]

Actually, I think my implementation has a potential issue: if one symbol is batched, and one symbol is not (like a weight matrix).

[MilesCranmer]

What do you think the correct strategy is here? Should I avoid adding einsum for keras, since it is technically a layer=True backend?

[arogozhnikov]

layer=True just refers to providing layers, it should not be related to any batch variables, and patterns should include batch variables. Anyway, I forgot keras now just a redirection to TF layers, so just excluded this part

[arogozhnikov]

PR is merged, made very minor changes.
thank you for paying attention to details and keeping pushing this!

[MilesCranmer]

Awesome!! Great to hear.

[gerdm]

This is a great PR!
+1 for “ Custom reduction operations”. Is there anyone already working on this?

[alok]

I'm interested in adding rearrange support. Any pointers?