Add C++ EinsumOp which support 2 operands einsum. (#42105) (#42357)

* full api fix

* when out is None, go old dygraph mode

* by static check

* first version: support 2-inputs forwards. TODO: 1. backward  2. BroadCast  3. MultiVariable

* time out -> 120

paddle/fluid/operators/einsum_op.cc

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <string>
+#include <vector>
+#include "paddle/fluid/framework/infershape_utils.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/phi/core/ddim.h"
+#include "paddle/phi/kernels/impl/einsum_impl.h"
+
+namespace paddle {
+namespace operators {
+class EinsumOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+};
+
+class EinsumOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("Operands", "(TensorList), The input tensor of einsum op.")
+        .AsDuplicable();
+    AddOutput("Out", "(Tensor), The output tensor of einsum op.");
+    AddAttr<std::string>("equation",
+                         "(string) A einsum equation. such as `ij,jk->ik`"
+                         "There must have `->` and the number of operands in "
+                         "equation must equals the `Operands` length.");
+    AddComment(R"DOC(
+Einsum Operator.
+
+This operator is used to perform einsum operation for given operands and equation.
+)DOC");
+  }
+};
+
+class EinsumGradOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    auto x_name = "Operands";
+    auto x_grad_name = framework::GradVarName(x_name);
+    ctx->SetOutputsDim(x_grad_name, ctx->GetInputsDim(x_name));
+    ctx->ShareAllLoD(x_name, x_grad_name);
+  }
+
+ protected:
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext& ctx) const override {
+    auto dtype = OperatorWithKernel::IndicateVarDataType(
+        ctx, framework::GradVarName("Out"));
+    return framework::OpKernelType(dtype, ctx.GetPlace());
+  }
+};
+
+template <typename T>
+class EinsumGradMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+  void Apply(GradOpPtr<T> retv) const override {
+    retv->SetType("einsum_grad");
+    retv->SetInput("Operands", this->Input("Operands"));
+    retv->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
+    retv->SetAttrMap(this->Attrs());
+    retv->SetOutput(framework::GradVarName("Operands"),
+                    this->InputGrad("Operands", false));
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+
+DECLARE_INFER_SHAPE_FUNCTOR(einsum, EinsumInferShapeFunctor,
+                            PD_INFER_META(phi::EinsumInferShape));
+
+REGISTER_OPERATOR(einsum, ops::EinsumOp, ops::EinsumOpMaker,
+                  EinsumInferShapeFunctor,
+                  ops::EinsumGradMaker<paddle::framework::OpDesc>,
+                  ops::EinsumGradMaker<paddle::imperative::OpBase>);
+
+REGISTER_OPERATOR(einsum_grad, ops::EinsumGradOp);

paddle/phi/kernels/cpu/einsum_grad_kernel.cc

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/einsum_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/einsum_grad_impl.h"
+
+PD_REGISTER_KERNEL(
+    einsum_grad, CPU, ALL_LAYOUT, phi::EinsumGradKernel, float, double) {}

paddle/phi/kernels/cpu/einsum_kernel.cc

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/einsum_kernel.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/einsum_impl.h"
+
+PD_REGISTER_KERNEL(einsum, CPU, ALL_LAYOUT, phi::EinsumKernel, float, double) {}

paddle/phi/kernels/einsum_grad_kernel.h

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void EinsumGradKernel(const Context& dev_ctx,
+                      const std::vector<const DenseTensor*>& x,
+                      const DenseTensor& out_grad,
+                      const std::string& equation,
+                      std::vector<DenseTensor*> x_grad);
+
+}  // namespace phi

paddle/phi/kernels/einsum_kernel.h

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void EinsumKernel(const Context& dev_ctx,
+                  const std::vector<const DenseTensor*>& inputs,
+                  const std::string& equation,
+                  DenseTensor* out);
+
+}  // namespace phi

paddle/phi/kernels/gpu/einsum_grad_kernel.cu

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/einsum_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/einsum_grad_impl.h"
+
+PD_REGISTER_KERNEL(
+    einsum_grad, GPU, ALL_LAYOUT, phi::EinsumGradKernel, float, double) {}

paddle/phi/kernels/gpu/einsum_kernel.cu

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/einsum_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/einsum_impl.h"
+
+PD_REGISTER_KERNEL(einsum, GPU, ALL_LAYOUT, phi::EinsumKernel, float, double) {}

paddle/phi/kernels/impl/einsum_grad_impl.h

+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/impl/einsum_impl.h"
+#include "paddle/phi/kernels/tile_kernel.h"
+#include "paddle/utils/string/string_helper.h"
+
+namespace phi {
+template <typename T, typename Context>
+DenseTensor PerformTileAndReduction(const Context& dev_ctx,
+                                    const LabelMap& label2type,
+                                    const LabelMap& label2shape,
+                                    const std::vector<int>& broadcast_dims,
+                                    const std::vector<int>& ellipsis_dims,
+                                    std::string op_label,  // value pass
+                                    DenseTensor& t) {      // NOLINT
+  ReplaceEllipsis(op_label);
+  DenseTensor ret;
+  std::vector<int> repeat_times;
+  std::vector<int> resize_dims;
+  std::vector<int> recover_shape;
+  for (int c : op_label) {
+    if (label2type[c] == LabelType::Reduction) {
+      // '.' can't be Reduction, so we don't deal '.' here.
+      repeat_times.push_back(label2shape[c]);
+      resize_dims.push_back(1);
+      recover_shape.push_back(label2shape[c]);
+    } else {
+      if (c != '.') {
+        resize_dims.push_back(label2shape[c]);
+        repeat_times.push_back(1);
+        recover_shape.push_back(label2shape[c]);
+      } else {
+        int n_dims = broadcast_dims.size();
+        resize_dims.insert(
+            resize_dims.end(), broadcast_dims.begin(), broadcast_dims.end());
+        recover_shape.insert(
+            recover_shape.end(), ellipsis_dims.begin(), ellipsis_dims.end());
+        while (n_dims--) repeat_times.push_back(1);
+      }
+    }
+  }
+  t.Resize(make_ddim(resize_dims));
+  DenseTensor after_tile;
+  TileKernel<T, Context>(dev_ctx, t, repeat_times, &after_tile);
+  size_t n_ellipsis_idx = op_label.find(".", 0);
+  if (n_ellipsis_idx != std::string::npos) {
+    // may be we need reduce. broadcast_dims is not equal to ellipsis dims.
+    std::vector<int64_t> to_reduce;
+    for (size_t i = 0; i < broadcast_dims.size() - ellipsis_dims.size(); ++i)
+      to_reduce.push_back(i + n_ellipsis_idx);
+
+    int new_offset =
+        n_ellipsis_idx + broadcast_dims.size() - ellipsis_dims.size();
+    for (size_t i = 0; i < ellipsis_dims.size(); ++i)
+      if (ellipsis_dims[i] == 1) to_reduce.push_back(i + new_offset);
+
+    VLOG(5) << "PermformTileAndReduction: reduce sum axis: "
+            << paddle::string::join_strings(to_reduce, ",");
+    if (to_reduce.size() != 0) {
+      ret = Sum<T, Context>(dev_ctx,
+                            after_tile,
+                            to_reduce,
+                            after_tile.dtype(),
+                            false);  // not keep dim.
+    } else {
+      ret = after_tile;
+    }
+  } else {
+    ret = after_tile;
+  }
+  VLOG(5) << "PermformTileAndReduction: recover shape: "
+          << paddle::string::join_strings(recover_shape, ",");
+  ret.Resize(make_ddim(recover_shape));
+  return ret;
+}
+
+template <typename T, typename Context>
+void EinsumGradKernel(const Context& dev_ctx,
+                      const std::vector<const DenseTensor*>& x,
+                      const DenseTensor& out_grad,
+                      const std::string& equation,
+                      std::vector<DenseTensor*> x_grad) {
+  VLOG(5) << "Start EisumGradKernel:";
+  LabelMap labelshape(0);
+  LabelMap labeltype(LabelType::Reduction);
+  std::vector<LabelMap> label2perms(x.size(), LabelMap(-1));
+  std::vector<char> all_labels;  // order: ABO, AO, BO, AB, Reduce
+  std::vector<std::vector<int>> ellipsis_dims(2);
+  std::vector<int> broadcast_dims;
+  std::vector<int> output_dims;
+
+  std::vector<DDim> input_dims;
+  for (auto& i : x) {
+    input_dims.push_back(i->dims());
+  }
+  std::string right;
+  ParseEinsumEquation(equation,
+                      input_dims,
+                      &labelshape,
+                      &labeltype,
+                      &all_labels,
+                      &label2perms,
+                      &ellipsis_dims,
+                      &broadcast_dims,
+                      &output_dims,
+                      &right);
+
+  auto gather_labels_except_reduction = [&labeltype](std::string all) {
+    std::string res("");
+    for (auto c : all)
+      if (labeltype[static_cast<int>(c)] != LabelType::Reduction) res += c;
+    return res;
+  };
+  if (x.size() == 1) {  // Unary
+    auto splits = paddle::string::split_string(equation, "->");
+    auto left = splits[0];
+    right = splits[1].substr(1);
+    auto new_equation = right + "->" + gather_labels_except_reduction(left);
+    auto new_operands = std::vector<const DenseTensor*>();
+    new_operands.push_back(&out_grad);
+    DenseTensor before_tile;
+    EinsumKernel<T, Context>(dev_ctx, new_operands, new_equation, &before_tile);
+    *(x_grad[0]) = PerformTileAndReduction<T, Context>(dev_ctx,
+                                                       labeltype,
+                                                       labelshape,
+                                                       broadcast_dims,
+                                                       ellipsis_dims[0],
+                                                       left,
+                                                       before_tile);
+  } else {
+    auto splits = paddle::string::split_string(equation, "->");
+    auto left = splits[0];
+    auto ops = paddle::string::split_string(left, ",");
+    right = splits[1].substr(1);
+
+    auto equation_for_A =
+        right + "," + ops[1] + "->" + gather_labels_except_reduction(ops[0]);
+    auto equation_for_B =
+        right + "," + ops[0] + "->" + gather_labels_except_reduction(ops[1]);
+    auto operands_for_A = std::vector<const DenseTensor*>();
+    auto operands_for_B = std::vector<const DenseTensor*>();
+    DenseTensor dA, dB;
+    operands_for_A.push_back(&out_grad);
+    operands_for_A.push_back(x[1]);
+    operands_for_B.push_back(&out_grad);
+    operands_for_B.push_back(x[0]);
+
+    DenseTensor before_tile;
+    EinsumKernel<T, Context>(dev_ctx, operands_for_A, equation_for_A, &dA);
+    EinsumKernel<T, Context>(dev_ctx, operands_for_B, equation_for_B, &dB);
+    *(x_grad[0]) = PerformTileAndReduction<T, Context>(dev_ctx,
+                                                       labeltype,
+                                                       labelshape,
+                                                       broadcast_dims,
+                                                       ellipsis_dims[0],
+                                                       ops[0],
+                                                       dA);
+    *(x_grad[1]) = PerformTileAndReduction<T, Context>(dev_ctx,
+                                                       labeltype,
+                                                       labelshape,
+                                                       broadcast_dims,
+                                                       ellipsis_dims[1],
+                                                       ops[1],
+                                                       dB);
+  }
+}
+}  // namespace phi

paddle/phi/ops/compat/einsum_sig.cc

+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/core/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature EinsumOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  return KernelSignature("einsum", {"Operands"}, {"equation"}, {"Out"});
+}
+
+KernelSignature EinsumGradOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  return KernelSignature("einsum_grad",
+                         {"Operands", {"Out@GRAD"}},
+                         {"equation"},
+                         {{"Operands@GRAD"}});
+}
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(einsum, phi::EinsumOpArgumentMapping);
+PD_REGISTER_ARG_MAPPING_FN(einsum_grad, phi::EinsumGradOpArgumentMapping);

python/paddle/fluid/tests/unittests/CMakeLists.txt

@@ -1063,6 +1063,7 @@ set_tests_properties(test_lstm_cudnn_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_stack_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_bilinear_interp_v2_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_svd_op PROPERTIES TIMEOUT 80)
+set_tests_properties(test_einsum_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_qr_op PROPERTIES TIMEOUT 60)
 set_tests_properties(test_deformable_psroi_pooling PROPERTIES TIMEOUT 120)
 set_tests_properties(test_trilinear_interp_v2_op PROPERTIES TIMEOUT 120)

python/paddle/fluid/tests/unittests/test_einsum_op.py

+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+import numpy as np
+import paddle
+from op_test import OpTest
+
+
+class TestEinsumBinary(OpTest):
+    def setUp(self):
+        paddle.enable_static()
+        self.op_type = "einsum"
+        self.disable = False
+        self.set_mandatory()
+        self.init_input()
+        np.random.seed(123)
+        out = np.einsum(self.equation, *self.inputs)
+        self.operands = []
+        for idx, inp in enumerate(self.inputs):
+            self.operands.append(("x" + str(idx), inp))
+        self.inputs = {"Operands": self.operands}
+        self.attrs = {"equation": self.equation}
+        self.outputs = {'Out': out}
+
+    def init_input(self):
+        self.inputs = []
+        for t, s in zip(self.types, self.shapes):
+            self.inputs.append(np.random.random(s).astype(t))
+
+    def set_mandatory(self):
+        self.disable = False
+        self.shapes = [(10, 10, 20), (20, 6)]
+        self.types = [np.float64, np.float64]
+        self.equation = "mij,jk->ki"
+
+    def test_check_output(self):
+        if not self.disable:
+            self.check_output()
+
+    def test_grad(self):
+        if not self.disable:
+            self.check_grad([op[0] for op in self.operands], ["Out"])
+
+
+class TestEinsum1(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(20, 3, 3), (20, 3, 3)]
+        self.types = [np.float64, np.float64]
+        self.equation = "mij,mjk->mik"
+
+
+class TestEinsum2(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(20, 3, 3), (20, 3, 3)]
+        self.types = [np.float64, np.float64]
+        self.equation = "mij,mjk->ikm"
+
+
+class TestEinsum3(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(10, 10), (10, 10)]
+        self.types = [np.float64, np.float64]
+        self.equation = "ij,jk->ik"  # }}}
+
+
+class TestEinsumWithReduction(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(10, 3, 5), (5, 30)]
+        self.types = [np.float64, np.float64]
+        self.equation = "ijk,kl->jl"
+
+
+class TestEinsumWithReduction1(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(10, 3, 3, 5), (10, 5, 10, 10)]
+        self.types = [np.float64, np.float64]
+        self.equation = "mijk,mklh->ljm"
+
+
+class TestEinsumWithUnary(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(10, 10, 3, 5)]
+        self.types = [np.float64]
+        self.equation = "mijk->mi"
+
+
+class TestEinsumWithUnary1(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(5, 10, 3, 3), (3, 6, 3, 10)]
+        self.types = [np.float64, np.float64]
+        self.equation = "imjl,jklm->imk"
+
+
+class TestEinsumWithBroadcast1(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(5, 10, 3, 3)]
+        self.types = [np.float64]
+        self.equation = "i...->..."
+
+
+class TestEinsumWithBroadcast2(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(10, 11), (3, 4, 5, 10)]
+        self.types = [np.float64, np.float64]
+        self.equation = "...ij,...i->j..."
+
+
+class TestEinsumWithBroadcast3(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(10, 3, 2, 3, 4), (12, 10)]
+        self.types = [np.float64, np.float64]
+        self.equation = "k...,...jk->...k"
+
+
+class TestEinsumWithBroadcast4(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(10, 3, 2, 3, 4), (12, 10)]
+        self.types = [np.float64, np.float64]
+        self.equation = "a...d,...cb->...abcd"
+
+
+class TestEinsumWithBroadcast5(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(3, 2, 2, 10), (10, 3, 2, 2)]
+        self.types = [np.float64, np.float64]
+        self.equation = "...a,a...->..."
+
+
+class TestEinsumWithBroadcast6(TestEinsumBinary):
+    def set_mandatory(self):
+        self.shapes = [(100), (100)]
+        self.types = [np.float64, np.float64]
+        self.equation = "i,i->"
+
+
+if __name__ == "__main__":
+    unittest.main()

python/paddle/tensor/einsum.py

@@ -20,6 +20,10 @@ from .linalg import dot, matmul, transpose
 from .manipulation import squeeze, unsqueeze, reshape
 from .math import multiply
 from .math import sum as paddle_sum
+from ..fluid.framework import _in_legacy_dygraph
+from paddle import _C_ops
+from ..fluid.data_feeder import check_variable_and_dtype, check_type, check_dtype
+from ..fluid.layer_helper import LayerHelper
 
 from paddle.common_ops_import import dygraph_only
 
@@ -660,6 +664,26 @@ def plan_einsum(operands, g_view, g_shape, g_supports, g_count, n_bcast):
     return plan
 
 
+def einsum_v2(equation, *operands):
+    if _in_legacy_dygraph():
+        # dygraph
+        return _C_ops.einsum(operands, 'equation', equation)
+    # static graph 
+    for inp in operands:
+        check_variable_and_dtype(inp, 'dtype', ['float32', 'float64'], 'einsum')
+    check_type(equation, 'equation', str, 'einsum')
+    helper = LayerHelper('einsum', **locals())
+    out = helper.create_variable_for_type_inference(dtype=operands[0].dtype)
+    attrs = dict()
+    attrs['equation'] = equation
+    helper.append_op(
+        type='einsum',
+        inputs={'Operands': operands},
+        outputs={'Out': out},
+        attrs=attrs, )
+    return out
+
+
 def einsum(equation, *operands):
     r"""
     einsum(equation, *operands)
@@ -817,6 +841,9 @@ def einsum(equation, *operands):
         #     [0.50226176, 0.24512935, 0.39881429],
         #     [0.51476848, 0.23367381, 0.39229113]]])
     """
+    import os
+    if int(os.environ.get('FLAGS_new_einsum', "0")):
+        return einsum_v2(equation, *operands)
 
     nop = len(operands)
     assert nop > 0, "At least one operand is expected."