Create einsum operation

einops/__init__.py

@@ -7,6 +7,7 @@ class EinopsError(RuntimeError):
     pass
 
 
-__all__ = ['rearrange', 'reduce', 'repeat', 'parse_shape', 'asnumpy', 'EinopsError']
+__all__ = ['rearrange', 'reduce', 'repeat', 'einsum',
+           'parse_shape', 'asnumpy', 'EinopsError']
 
-from .einops import rearrange, reduce, repeat, parse_shape, asnumpy
+from .einops import rearrange, reduce, repeat, einsum, parse_shape, asnumpy

einops/_backends.py

@@ -12,6 +12,7 @@
 
 import sys
 import warnings
+import functools
 
 __author__ = 'Alex Rogozhnikov'
 
@@ -117,6 +118,9 @@ def layers(self):
     def __repr__(self):
         return "<einops backend for {}>".format(self.framework_name)
 
+    def einsum(self, pattern, *x):
+        raise NotImplementedError("backend does not support einsum")
+
 
 class UnknownSize:
     """ pseudo-symbol for symbolic frameworks which do not provide symbols for shape elements """
@@ -168,6 +172,9 @@ def is_float_type(self, x):
     def add_axis(self, x, new_position):
         return self.np.expand_dims(x, new_position)
 
+    def einsum(self, pattern, *x):
+        return self.np.einsum(pattern, *x)
+
 
 class JaxBackend(NumpyBackend):
     framework_name = 'jax'
@@ -353,6 +360,9 @@ def layers(self):
         from .layers import torch
         return torch
 
+    def einsum(self, pattern, *x):
+        return self.torch.einsum(pattern, *x)
+
 
 class CupyBackend(AbstractBackend):
     framework_name = 'cupy'
@@ -385,6 +395,9 @@ def add_axis(self, x, new_position):
     def is_float_type(self, x):
         return x.dtype in ('float16', 'float32', 'float64', 'float128')
 
+    def einsum(self, pattern, *x):
+        return self.cupy.einsum(pattern, *x)
+
 
 class ChainerBackend(AbstractBackend):
     framework_name = 'chainer'
@@ -428,6 +441,9 @@ def layers(self):
         from .layers import chainer
         return chainer
 
+    def einsum(self, pattern, *x):
+        return self.chainer.functions.einsum(pattern, *x)
+
 
 class HashableTuple:
     """Overcomes non-hashability of symbolic elements"""
@@ -507,6 +523,9 @@ def layers(self):
         from .layers import tensorflow
         return tensorflow
 
+    def einsum(self, pattern, *x):
+        return self.tf.einsum(pattern, *x)
+
 
 class KerasBackend(AbstractBackend):
     framework_name = 'tensorflow.keras'
@@ -524,9 +543,8 @@ def create_symbol(self, shape):
         return self.keras.Input(batch_shape=shape)
 
     def eval_symbol(self, symbol, input_dict):
-        (variable, value), = input_dict
-        model = self.keras.models.Model(variable, symbol)
-        return model.predict_on_batch(value)
+        model = self.keras.models.Model([var for (var, _) in input_dict], symbol)
+        return model.predict_on_batch([val for (_, val) in input_dict])
 
     def arange(self, start, stop):
         return self.K.arange(start, stop)
@@ -560,6 +578,12 @@ def layers(self):
         from .layers import keras
         return keras
 
+    def einsum(self, pattern, *x):
+        return self.tf.vectorized_map(
+            functools.partial(self.tf.einsum, pattern),
+            *x
+        )
+
 
 class OneFlowBackend(AbstractBackend):
     framework_name = "oneflow"
@@ -621,4 +645,7 @@ def is_float_type(self, x):
 
     def layers(self):
         from .layers import oneflow
-        return oneflow
+        return oneflow
+
+    def einsum(self, pattern, *x):
+        return self.flow.einsum(pattern, *x)

einops/einops.py

@@ -1,5 +1,6 @@
 import functools
 import itertools
+import string
 import typing
 from collections import OrderedDict
 from typing import Tuple, List, Dict, Union, Callable, Optional, TypeVar
@@ -624,3 +625,144 @@ def asnumpy(tensor) -> 'numpy.ndarray':
         `numpy.ndarray`, converted to numpy
     """
     return get_backend(tensor).to_numpy(tensor)
+
+def _validate_einsum_axis_name(axis_name):
+    if len(axis_name) == 0:
+        raise NotImplementedError("Singleton () axes are not yet supported in einsum.")
+    if len(axis_name) > 1:
+        raise NotImplementedError("Shape rearrangement is not yet supported in einsum.")
+
+    axis_name = axis_name[0]
+
+    if len(axis_name) == 0:
+        raise RuntimeError("Encountered empty axis name in einsum.")
+    if not isinstance(axis_name, str):
+        raise RuntimeError("Axis name in einsum must be a string.")
+    if axis_name[0].isdigit():
+        raise RuntimeError("Axis name in einsum must not start with a number.")
+
+
+@functools.lru_cache(256)
+def _compactify_pattern_for_einsum(pattern: str) -> str:
+    if "->" not in pattern:
+        # numpy allows this, so make sure users
+        # don't accidentally do something like this.
+        raise ValueError("Einsum pattern must contain '->'.")
+    lefts, right = pattern.split('->')
+    lefts = lefts.split(',')
+
+    lefts = [
+        ParsedExpression(left, allow_underscore=True, allow_duplicates=True)
+        for left in lefts
+    ]
+
+    right = ParsedExpression(right, allow_underscore=True)
+
+    # Start from a, and go up to Z
+    output_axis_names = string.ascii_letters
+    i = 0
+    axis_name_mapping = {}
+
+    left_patterns = []
+    for left in lefts:
+        left_pattern = ""
+        for raw_axis_name in left.composition:
+
+            if raw_axis_name == _ellipsis:
+                left_pattern += '...'
+                continue
+
+            _validate_einsum_axis_name(raw_axis_name)
+            axis_name = raw_axis_name[0]
+            if axis_name not in axis_name_mapping:
+                if i >= len(output_axis_names):
+                    raise RuntimeError("Too many axes in einsum.")
+                axis_name_mapping[axis_name] = output_axis_names[i]
+                i += 1
+
+            left_pattern += axis_name_mapping[axis_name]
+        left_patterns.append(left_pattern)
+
+    output_pattern = ",".join(left_patterns) + "->"
+
+    for raw_axis_name in right.composition:
+        if raw_axis_name == _ellipsis:
+            output_pattern += '...'
+            continue
+
+        _validate_einsum_axis_name(raw_axis_name)
+        axis_name = raw_axis_name[0]
+
+        if axis_name not in axis_name_mapping:
+            raise RuntimeError("Unknown axis on right side of einsum.")
+
+        output_pattern += axis_name_mapping[axis_name]
+
+    return output_pattern
+
+
+def einsum(pattern: str, *tensors: List[Tensor]) -> Tensor:
+    """
+    einops.einsum calls einsum operations with einops-style named
+    axes indexing, computing tensor products with an arbitrary
+    number of tensors.
+
+    Note that unlike other einops functions, here you must give
+    the pattern before the tensor(s), rather than after.
+    Also, note that rearrange operations such as `"(batch chan) out"`,
+    or singleton axes `()`, are not currently supported.
+
+    Examples:
+
+    For a given pattern such as:
+    ```python
+    >>> x, y, z = np.random.randn(3, 20, 20, 20)
+    >>> output = einsum("a b c, c b d, a g k -> a b k", x, y, z)
+
+    ```
+    the following formula is computed:
+    ```tex
+    output[a, b, k] = 
+        \sum_{c, d, g} x[a, b, c] * y[c, b, d] * z[a, g, k]
+    ```
+    where the summation over `c`, `d`, and `g` is performed
+    because those axes names do not appear on the right-hand side.
+
+    Let's see some additional examples:
+    ```python
+    # Filter a set of images:
+    >>> batched_images = np.random.randn(128, 16, 16)
+    >>> filters = np.random.randn(16, 16, 30)
+    >>> result = einsum("batch h w, h w channel -> batch channel",
+    ...                 batched_images, filters) 
+    >>> 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("out_dim in_dim, ... in_dim -> ... out_dim",
+    ...                 weights, data)
+    >>> result.shape
+    (50, 30, 10)
+
+    # Matrix trace on a single tensor:
+    >>> matrix = np.random.randn(10, 10)
+    >>> result = einsum("i i ->", matrix)
+    >>> result.shape
+    ()
+
+    ```
+
+    Parameters:
+        pattern: string, einsum pattern, with commas
+                 separating specifications for each tensor.
+        tensors: tensors of any supported library (numpy, tensorflow, pytorch, jax).
+
+    Returns:
+        Tensor of the same type as input, after processing with einsum.
+
+    """
+    pattern = _compactify_pattern_for_einsum(pattern)
+    return get_backend(tensors[0]).einsum(pattern, *tensors)

einops/parsing.py

@@ -26,7 +26,8 @@ class ParsedExpression:
     non-mutable structure that contains information about one side of expression (e.g. 'b c (h w)')
     and keeps some information important for downstream
     """
-    def __init__(self, expression, *, allow_underscore: bool = False):
+    def __init__(self, expression, *, allow_underscore: bool = False,
+                 allow_duplicates: bool = False):
         self.has_ellipsis: bool = False
         self.has_ellipsis_parenthesized: Optional[bool] = None
         self.identifiers: Set[str] = set()
@@ -48,7 +49,7 @@ def __init__(self, expression, *, allow_underscore: bool = False):
         def add_axis_name(x):
             if x is not None:
                 if x in self.identifiers:
-                    if not (allow_underscore and x == "_"):
+                    if not (allow_underscore and x == "_") and not allow_duplicates:
                         raise EinopsError('Indexing expression contains duplicate dimension "{}"'.format(x))
                 if x == _ellipsis:
                     self.identifiers.add(_ellipsis)