Add Einsum

Author: Laicheng0830

docs/modules/ops.rst

@@ -116,6 +116,7 @@ API - Operations
    diag
    mask_select
    eye
+   einsum
 
 TensorLayerX Tensor Operations
 --------------------------------
@@ -558,4 +559,8 @@ mask_select
 
 eye
 ^^^^^^^^^^^^^^^^^^^^^^^
-.. autofunction:: eye
+.. autofunction:: eye
+
+einsum
+^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: einsum

examples/basic_tutorials/tutorial_tensorlayerx_graph.py

@@ -1,6 +1,11 @@
 #! /usr/bin/python
 # -*- coding: utf-8 -*-
 
+import os
+# os.environ['TL_BACKEND'] = 'tensorflow'
+# os.environ['TL_BACKEND'] = 'mindspore'
+os.environ['TL_BACKEND'] = 'torch'
+
 import tensorlayerx as tlx
 from tensorlayerx.nn import Module
 from tensorlayerx.nn import Linear, Conv2d, BatchNorm2d, MaxPool2d, Flatten

tensorlayerx/backend/ops/__init__.py

@@ -84,6 +84,7 @@
 from .load_backend import QuanConvBn
 
 # load ops
+from .load_backend import einsum
 from .load_backend import Variable
 from .load_backend import matmul
 from .load_backend import add
@@ -245,3 +246,4 @@
 from .load_backend import ClipGradByValue
 from .load_backend import ClipGradByNorm
 from .load_backend import ClipByGlobalNorm
+from .load_backend import Einsum

tensorlayerx/backend/ops/mindspore_backend.py

@@ -1821,3 +1821,21 @@ def eye(n, m = None, dtype = None):
     if dtype is None:
         dtype = mstype.float32
     return ms.numpy.eye(n, m, dtype = dtype)
+
+
+def einsum(equation, *operands):
+    if ms.__version__ < '1.7.0':
+        raise NotImplementedError("Only MindSpore versions later than 1.7.0 are supported.")
+    einsum_obj = ms.ops.Einsum(equation)
+    return einsum_obj(tuple(operands))
+
+
+class Einsum(Cell):
+    def __init__(self, equation):
+        super(Einsum, self).__init__()
+        if ms.__version__ < '1.7.0':
+            raise NotImplementedError("Only MindSpore versions later than 1.7.0 are supported.")
+        self.einsum = ms.ops.Einsum(equation)
+
+    def __call__(self, *args):
+        return self.einsum(tuple(args))

tensorlayerx/backend/ops/paddle_backend.py

@@ -1834,3 +1834,25 @@ def _apply_mask_1d(reshaped_tensor, mask, axis=None):
 
 def eye(n, m=None, dtype=None):
     return paddle.eye(n, m, dtype)
+
+
+def einsum(equation, *operands):
+    try:
+        from paddlenlp.ops import einsum
+    except:
+        raise Exception("Paddlenlp needs to be installed.")
+    return einsum(equation, *operands)
+
+
+class Einsum(object):
+    def __init__(self, equation):
+        super(Einsum, self).__init__()
+        try:
+            from paddlenlp.ops import einsum
+        except:
+            raise Exception("Paddlenlp needs to be installed.")
+        self.equation = equation
+
+    def __call__(self, *args):
+        return einsum(self.equation, *args)
+

tensorlayerx/backend/ops/tensorflow_backend.py

@@ -3827,3 +3827,44 @@ def eye(n, m = None, dtype = None):
         dtype = tf.dtypes.float32
     return tf.eye(n, m, dtype = dtype)
 
+
+def einsum(equation, *operands):
+    """
+    Sums the product of the elements of the input operands along dimensions specified
+    using a notation based on the Einstein summation convention.
+
+    Parameters
+    ----------
+    equation : An attribute
+        represent the operation you want to do.
+        the value can contain only letters([a-z][A-Z]), commas(,), ellipsis(…), and arrow(->).
+        the letters represent inputs’s tensor dimension, commas(,)represent separate tensors, ellipsis(…) indicates
+        the tensor dimension that you do not care about, the left of the arrow(->) indicates the input tensors,
+        and the right of it indicates the desired output dimension.
+
+    operands : list
+        input tensor used for calculation. the data type of the tensor must be the same.
+
+    Returns
+    -------
+        Tensor, the shape of it can be obtained from the equation, and the data type is the same as input tensors.
+
+    Examples
+    ---------
+    >>> import tensorlayerx as tlx
+    >>> x = tlx.nn.Input((5,))
+    >>> y = tlx.nn.Input((4,))
+    >>> out = tlx.ops.einsum('i,j->ij', x, y)
+    >>> cal_enisum = tlx.ops.Einsum('i,j->ij')
+    >>> out = cal_enisum(x, y)
+    """
+    return tf.einsum(equation, *operands)
+
+
+class Einsum(object):
+    def __init__(self, equation):
+        super(Einsum, self).__init__()
+        self.equation = equation
+
+    def __call__(self, *args):
+        return tf.einsum(self.equation, *args)

tensorlayerx/backend/ops/torch_backend.py

@@ -1666,4 +1666,17 @@ def mask_select(x, mask, axis = 0):
         return x[:,:,:, mask]
 
 def eye(n, m=None, dtype=None):
-    return torch.eye(n = n, m = m, dtype =dtype)
+    return torch.eye(n = n, m = m, dtype =dtype)
+
+
+def einsum(equation, *operands):
+    return torch.einsum(equation, *operands)
+
+
+class Einsum(object):
+    def __init__(self, equation):
+        super(Einsum, self).__init__()
+        self.equation = equation
+
+    def __call__(self, *args):
+        return torch.einsum(self.equation, *args)

tensorlayerx/nn/core/core_mindspore.py

@@ -287,8 +287,7 @@ def init_build(self, *inputs, **kwargs):
 
     def build_graph(self, *inputs, **kwargs):
         # Add nodes only when the composition is needed.
-        layers = self.cells_and_names(name_prefix='')
-        for layer_name, layer in layers:
+        for layer_name, layer in self._cells.items():
             if isinstance(layer, Module):
                 layer._build_graph = True
         self.set_eval()

tensorlayerx/nn/core/core_torch.py

@@ -174,7 +174,7 @@ def init_build(self, *inputs, **kwargs):
 
     def build_graph(self, *inputs, **kwargs):
         # Add nodes only when the composition is needed.
-        for name, layer in self.named_modules():
+        for name, layer in self._modules.items():
             if isinstance(layer, Module):
                 layer._build_graph = True
         self.set_eval()