From 5bcfb82ec6e6172af50774e8608624d8bdd3324b Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Thu, 30 Jul 2020 14:38:18 +0800
Subject: [PATCH 01/14] [TOPI] Einsum * Fix tuple * fix oshape

---
 3q                           |  28 ++
 include/tvm/topi/transform.h | 827 ++++++++++++++++++++++++++++++++++-
 python/tvm/topi/transform.py |   7 +
 src/topi/transform.cc        |   8 +
 4 files changed, 868 insertions(+), 2 deletions(-)
 create mode 100644 3q

diff --git a/3q b/3q
new file mode 100644
index 000000000000..57d1462d5996
--- /dev/null
+++ b/3q
@@ -0,0 +1,28 @@
+pick 6a175f57a * debug
+s caf7f5e64 * d
+
+# Rebase 4b61a9764..caf7f5e64 onto 4b61a9764 (2 commands)
+#
+# Commands:
+# p, pick <commit> = use commit
+# r, reword <commit> = use commit, but edit the commit message
+# e, edit <commit> = use commit, but stop for amending
+# s, squash <commit> = use commit, but meld into previous commit
+# f, fixup <commit> = like "squash", but discard this commit's log message
+# x, exec <command> = run command (the rest of the line) using shell
+# b, break = stop here (continue rebase later with 'git rebase --continue')
+# d, drop <commit> = remove commit
+# l, label <label> = label current HEAD with a name
+# t, reset <label> = reset HEAD to a label
+# m, merge [-C <commit> | -c <commit>] <label> [# <oneline>]
+# .       create a merge commit using the original merge commit's
+# .       message (or the oneline, if no original merge commit was
+# .       specified). Use -c <commit> to reword the commit message.
+#
+# These lines can be re-ordered; they are executed from top to bottom.
+#
+# If you remove a line here THAT COMMIT WILL BE LOST.
+#
+# However, if you remove everything, the rebase will be aborted.
+#
+# Note that empty commits are commented out
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index 261fdf9970a3..94f1c50f26ec 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -23,6 +23,9 @@
  */
 #ifndef TVM_TOPI_TRANSFORM_H_
 #define TVM_TOPI_TRANSFORM_H_
+#define NPY_MAXDIMS 16
+#define NPY_MAXARGS 16
+#define MAXAXIS 128
 
 #include <tvm/te/operation.h>
 #include <tvm/tir/data_layout.h>
@@ -38,6 +41,7 @@
 #include <string>
 #include <unordered_set>
 #include <vector>
+#include <bitset>
 
 #include "detail/broadcast.h"
 
@@ -1266,8 +1270,8 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, int axes = 2,
   Array<IterVar> iter_vars;
   for (int i = 0; i < axes; ++i)
     iter_vars.push_back(reduce_axis(Range(0, B->shape[i]), "k" + std::to_string(i)));
-
   auto func = [&A, &B, &iter_vars, axes](const Array<Var>& input_indices) {
+
     Array<PrimExpr> A_indices(input_indices.begin(),
                               input_indices.begin() + (A->shape.size() - axes));
     for (auto& v : iter_vars) A_indices.push_back(v);
@@ -1281,8 +1285,9 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, int axes = 2,
     // Some passes don't like reductions with empty axis, so avoid it here
     if (iter_vars.empty())
       return A(A_indices) * B(B_indices);
-    else
+    else {
       return sum(A(A_indices) * B(B_indices), iter_vars);
+    }
   };
 
   return compute(output_shape, func, name, tag);
@@ -1344,6 +1349,824 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, Array<PrimExp
   return compute(output_shape, func, name, tag);
 }
 
+inline Array<PrimExpr> get_stride(const Array<PrimExpr> shape) {
+  size_t ndim = shape.size();
+  int prod = 1;
+  Array<PrimExpr> stride = Array<PrimExpr>(ndim, -1);
+  for (int i = ndim - 1; i >= 0; i--) {
+    stride.Set(i, if_then_else(shape[i] > 1, prod, 0));
+    prod = prod * GetConstInt(shape[i]);
+  }
+  return stride;
+}
+
+inline Array<PrimExpr> pad(const Array<PrimExpr> shape, int odim) {
+  int ndim = shape.size();
+  CHECK_GE(odim, ndim);
+  Array<PrimExpr> ret(static_cast<size_t>(odim), 1);
+  for (int idim = 0; idim < ndim; ++idim) {
+    ret.Set(idim, shape[idim]);
+  }
+  return ret;
+}
+
+inline int parse_operand_subscripts(const char *subscripts, int length,
+                                    int ndim, int iop, char *op_labels,
+                                    char *label_counts, int *min_label, int *max_label) {
+  int i;
+  int idim = 0;
+  int ellipsis = -1;
+
+  /* Process all labels for this operand */
+  for (i = 0; i < length; ++i) {
+    int label = subscripts[i];
+
+    /* A proper label for an axis. */
+    if (label > 0 && isalpha(label)) {
+      /* Check we don't exceed the operator dimensions. */
+      CHECK(idim < ndim)
+        << "einstein sum subscripts string contains "
+        << "too many subscripts for operand "
+        << iop;
+
+      op_labels[idim++] = label;
+      if (label < *min_label) {
+        *min_label = label;
+      }
+      if (label > *max_label) {
+        *max_label = label;
+      }
+      label_counts[label]++;
+    } else if (label == '.') {
+      /* The beginning of the ellipsis. */
+      /* Check it's a proper ellipsis. */
+      CHECK(!(ellipsis != -1 || i + 2 >= length
+              || subscripts[++i] != '.' || subscripts[++i] != '.'))
+        << "einstein sum subscripts string contains a "
+        << "'.' that is not part of an ellipsis ('...') "
+        << "in operand "
+        << iop;
+
+      ellipsis = idim;
+    } else {
+        CHECK(label == ' ')
+          << "invalid subscript '" << static_cast<char>(label)
+          << "' in einstein sum "
+          << "subscripts string, subscripts must "
+          << "be letters";
+    }
+  }
+
+  /* No ellipsis found, labels must match dimensions exactly. */
+  if (ellipsis == -1) {
+    CHECK(idim == ndim)
+      << "operand has more dimensions than subscripts "
+      << "given in einstein sum, but no '...' ellipsis "
+      << "provided to broadcast the extra dimensions.";
+  } else if (idim < ndim) {
+    /* Ellipsis found, may have to add broadcast dimensions. */
+    /* Move labels after ellipsis to the end. */
+    for (i = 0; i < idim - ellipsis; ++i) {
+      op_labels[ndim - i - 1] = op_labels[idim - i - 1];
+    }
+    /* Set all broadcast dimensions to zero. */
+    for (i = 0; i < ndim - idim; ++i) {
+      op_labels[ellipsis + i] = 0;
+    }
+  }
+
+  /*
+   * Find any labels duplicated for this operand, and turn them
+   * into negative offsets to the axis to merge with.
+   *
+   * In C, the char type may be signed or unsigned, but with
+   * twos complement arithmetic the char is ok either way here, and
+   * later where it matters the char is cast to a signed char.
+   */
+  for (idim = 0; idim < ndim - 1; ++idim) {
+    int label = op_labels[idim];
+    /* If it is a proper label, find any duplicates of it. */
+    if (label > 0) {
+      /* Search for the next matching label. */
+      char *next = reinterpret_cast<char*>(memchr(op_labels + idim + 1, label, ndim - idim - 1));
+
+      while (next != nullptr) {
+        /* The offset from next to op_labels[idim] (negative). */
+        *next = static_cast<char>((op_labels + idim) - next);
+        /* Search for the next matching label. */
+        next = reinterpret_cast<char*>(memchr(next + 1, label, op_labels + ndim - 1 - next));
+      }
+    }
+  }
+  return 0;
+}
+
+inline int parse_output_subscripts(const char *subscripts, int length,
+                                   int ndim_broadcast,
+                                   const char *label_counts, char *out_labels) {
+  int i, bdim;
+  int ndim = 0;
+  int ellipsis = 0;
+
+  /* Process all the output labels. */
+  for (i = 0; i < length; ++i) {
+    int label = subscripts[i];
+
+    /* A proper label for an axis. */
+    if (label > 0 && isalpha(label)) {
+      /* Check that it doesn't occur again. */
+      CHECK(memchr(subscripts + i + 1, label, length - i - 1) == nullptr)
+        << "einstein sum subscripts string includes "
+        << "output subscript '" << static_cast<char>(label)
+        << "' multiple times";
+
+      /* Check that it was used in the inputs. */
+      CHECK(label_counts[label] != 0)
+        << "einstein sum subscripts string included "
+        << "output subscript '" << static_cast<char>(label)
+        << "' which never appeared "
+        << "in an input";
+
+      /* Check that there is room in out_labels for this label. */
+      CHECK(ndim < NPY_MAXDIMS)
+        << "einstein sum subscripts string contains "
+        << "too many subscripts in the output";
+
+      out_labels[ndim++] = label;
+    } else if (label == '.') {
+      /* The beginning of the ellipsis. */
+      /* Check it is a proper ellipsis. */
+      CHECK(!(ellipsis || i + 2 >= length
+              || subscripts[++i] != '.' || subscripts[++i] != '.'))
+        << "einstein sum subscripts string "
+        << "contains a '.' that is not part of "
+        << "an ellipsis ('...') in the output";
+
+      /* Check there is room in out_labels for broadcast dims. */
+      CHECK(ndim + ndim_broadcast <= NPY_MAXDIMS)
+        << "einstein sum subscripts string contains "
+        << "too many subscripts in the output";
+
+      ellipsis = 1;
+      for (bdim = 0; bdim < ndim_broadcast; ++bdim) {
+        out_labels[ndim++] = 0;
+      }
+    } else {
+      CHECK(label == ' ')
+        << "invalid subscript '" << static_cast<char>(label)
+        << "' in einstein sum "
+        << "subscripts string, subscripts must "
+        << "be letters";
+    }
+  }
+
+  /* If no ellipsis was found there should be no broadcast dimensions. */
+  CHECK(!(!ellipsis && ndim_broadcast > 0))
+    << "output has more dimensions than subscripts "
+    << "given in einstein sum, but no '...' ellipsis "
+    << "provided to broadcast the extra dimensions.";
+
+  return ndim;
+}
+
+inline void get_combined_dims_view(const Tensor& op, int iop,
+                                   char *labels,
+                                   Array<PrimExpr> newshape,
+                                   Array<PrimExpr> newstride) {
+  int idim, ndim, icombine, combineoffset;
+  int icombinemap[NPY_MAXDIMS];
+  int newdim;
+
+  Array<PrimExpr> shape = op->shape;
+  Array<PrimExpr> stride = get_stride(shape);
+  ndim = op.ndim();
+  newdim = newshape.size();
+
+  /* Initialize the dimensions and strides to zero */
+  for (idim = 0; idim < newdim; ++idim) {
+    newshape.Set(idim, 0);
+    newstride.Set(idim, 0);
+  }
+
+  /* Copy the dimensions and strides, except when collapsing */
+  icombine = 0;
+  for (idim = 0; idim < ndim; ++idim) {
+    /*
+     * The char type may be either signed or unsigned, we
+     * need it to be signed here.
+     */
+    int label = (signed char)labels[idim];
+    /* If this label says to merge axes, get the actual label */
+    if (label < 0) {
+      combineoffset = label;
+      label = labels[idim+label];
+    } else {
+      combineoffset = 0;
+      if (icombine != idim) {
+        labels[icombine] = labels[idim];
+      }
+      icombinemap[idim] = icombine;
+    }
+    /* If the label is 0, it's an unlabeled broadcast dimension */
+    if (label == 0) {
+      newshape.Set(icombine, shape[idim]);
+      newstride.Set(icombine, stride[idim]);
+    } else {
+      /* Update the combined axis dimensions and strides */
+      int i = icombinemap[idim + combineoffset]; 
+      CHECK(!(
+        (combineoffset < 0) && GetConstInt(newshape[i] != 0 && newshape[i] != shape[idim])
+        ))
+        << "dimensions in operand " << iop
+        << " for collapsing index '" << label
+        << "' don't match (" << GetConstInt(newshape[i])
+        << " != " << shape[idim] << ")";
+      newshape.Set(i, shape[idim]);
+      newstride.Set(i, newstride[i] + stride[idim]);
+    }
+
+    /* If the label didn't say to combine axes, increment dest i */
+    if (combineoffset == 0) {
+      icombine++;
+    }
+  }
+}
+
+inline static int prepare_op_axes(int ndim, int iop, char *labels,
+                                  int *axes, int ndim_iter, char *iter_labels) {
+  int i, label, ibroadcast;
+
+  ibroadcast = ndim-1;
+  for (i = ndim_iter-1; i >= 0; --i) {
+    label = iter_labels[i];
+    /*
+     * If it's an unlabeled broadcast dimension, choose
+     * the next broadcast dimension from the operand.
+     */
+    if (label == 0) {
+      while (ibroadcast >= 0 && labels[ibroadcast] != 0) {
+        --ibroadcast;
+      }
+      /*
+       * If we used up all the operand broadcast dimensions,
+       * extend it with a "newaxis"
+       */
+      if (ibroadcast < 0) {
+        axes[i] = -1;
+      } else {
+        /* Otherwise map to the broadcast axis */
+        axes[i] = ibroadcast;
+        --ibroadcast;
+      }
+    } else {
+      /* It's a labeled dimension, find the matching one */
+      char *match = reinterpret_cast<char*>(memchr(labels, label, ndim));
+      /* If the op doesn't have the label, broadcast it */
+      if (match == nullptr) {
+        axes[i] = -1;
+      } else {
+        /* Otherwise use it */
+        axes[i] = match - labels;
+      }
+    }
+  }
+  return 0;
+}
+
+inline int _count_substring(const std::string& str,
+                            const std::string& sub) {
+  int count = 0;
+  std::string::size_type pos = 0;
+  while ((pos = str.find(sub, pos)) != std::string::npos) {
+    ++count;
+    pos += sub.length();
+  }
+  return count;
+}
+
+inline std::bitset<MAXAXIS> str2set(const std::string& str) {
+  std::bitset<MAXAXIS> ret;
+  for (const char& c : str) {
+    ret.set(static_cast<int>(c));
+  }
+  return ret;
+}
+
+inline std::vector<std::string> split(const std::string& str,
+                                      const std::string& sub) {
+  std::string::size_type pos = 0;
+  std::string::size_type start = 0;
+  std::vector<std::string> ret;
+  while ((pos = str.find(sub, start)) != std::string::npos) {
+    ret.push_back(str.substr(start, pos - start));
+    start = pos + sub.length();
+  }
+  ret.push_back(str.substr(start));
+  return ret;
+}
+
+inline std::vector<std::string> _parse_einsum_input(
+  std::string subscripts,
+  const std::vector<Array<PrimExpr>>& operands) {
+  const std::string einsum_symbols =
+    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
+  std::bitset<MAXAXIS> einsum_symbols_set;
+  for (const char& c : einsum_symbols) {
+    einsum_symbols_set.set(c);
+  }
+
+  CHECK_NE(operands.size(), 0U)
+    << "No input operands";
+
+  auto end_pos = std::remove(subscripts.begin(), subscripts.end(), ' ');
+  subscripts.erase(end_pos, subscripts.end());
+
+  // Ensure all characters are valid
+  for (const char& c : subscripts) {
+    if (c == '.' || c == ',' || c == '-' || c == '>') {
+      continue;
+    }
+    CHECK(einsum_symbols_set.test(c))
+      << "Character " << c
+      << " is not a valid symbol.";
+  }
+
+  // Check for proper "->"
+  if (subscripts.find('-') != std::string::npos ||
+      subscripts.find('>') != std::string::npos) {
+    bool invalid = (std::count(subscripts.begin(), subscripts.end(), '-') > 1 ||
+                    std::count(subscripts.begin(), subscripts.end(), '>') > 1);
+    CHECK(!invalid && _count_substring(subscripts, "->") == 1)
+      << "Subscripts can only contain one '->'.";
+  }
+
+  // Parse ellipses
+  if (subscripts.find('.') != std::string::npos) {
+    std::string used = subscripts;
+    used.erase(std::remove_if(used.begin(),
+                              used.end(),
+                              [](const char& c){return c == '.' ||
+                                                       c == ',' ||
+                                                       c == '-' ||
+                                                       c == '>';}),
+               used.end());
+
+    std::bitset<MAXAXIS> used_set = str2set(used);
+    std::string ellipse_inds = "";
+    for (const char& c : einsum_symbols) {
+      if (!used_set.test(static_cast<int>(c))) {
+        ellipse_inds.append(1, c);
+      }
+    }
+    int longest = 0;
+    std::string input_tmp, output_sub;
+    std::vector<std::string> split_subscripts;
+    bool out_sub;
+
+    if (subscripts.find("->") != std::string::npos) {
+      std::vector<std::string> tmp = split(subscripts, "->");
+      input_tmp = tmp[0];
+      output_sub = tmp[1];
+      split_subscripts = split(input_tmp, ",");
+      out_sub = true;
+    } else {
+      split_subscripts = split(subscripts, ",");
+      out_sub = false;
+    }
+
+    size_t size_split_subscripts = split_subscripts.size();
+    subscripts = "";
+    for (size_t i = 0; i < size_split_subscripts; ++i) {
+      const std::string& sub = split_subscripts[i];
+      if (sub.find('.') != std::string::npos) {
+        CHECK_EQ(std::count(sub.begin(), sub.end(), '.'), 3)
+          << "Invalid Ellipses";
+        CHECK_EQ(_count_substring(sub, "..."), 1)
+          << "Invalid Ellipses";
+
+        // Take into account numerical values
+        int ellipse_count = 0;
+        if (operands[i].size() == 0) {
+          ellipse_count = 0;
+        } else {
+          ellipse_count = std::max(operands[i].size(), static_cast<size_t>(1));
+          ellipse_count -= sub.length() - 3;
+        }
+
+        if (ellipse_count > longest) {
+          longest = ellipse_count;
+        }
+
+        CHECK_GE(ellipse_count, 0)
+          << "Ellipses lengths do not match.";
+        if (ellipse_count == 0) {
+          split_subscripts[i].erase(sub.find("..."), 3);
+        } else {
+          std::string rep_inds = ellipse_inds.substr(ellipse_inds.length() - ellipse_count);
+          split_subscripts[i].replace(sub.find("..."), 3, rep_inds);
+        }
+      }
+      subscripts += split_subscripts[i];
+      if (i + 1 < size_split_subscripts) {
+        subscripts += ",";
+      }
+    }
+    std::string out_ellipse;
+    if (longest == 0) {
+      out_ellipse = "";
+    } else {
+      out_ellipse = ellipse_inds.substr(ellipse_inds.length() - longest);
+    }
+
+    if (out_sub) {
+      output_sub.replace(output_sub.find("..."), 3, out_ellipse);
+      subscripts += "->" + output_sub;
+    } else {
+      // Special care for outputless ellipses
+      std::bitset<MAXAXIS> out_ellipse_set = str2set(out_ellipse);
+      std::string tmp_subscripts = subscripts, output_subscript = "";
+      size_t len_tmp_subscripts = tmp_subscripts.length();
+      std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
+      for (size_t i = 0; i < len_tmp_subscripts; ++i) {
+        const char& c = tmp_subscripts[i];
+        if (c == ',') {
+          continue;
+        }
+        CHECK(einsum_symbols_set.test(c))
+          << "Character " << c
+          << " is not a valid symbol.";
+        if ((i == 0 || tmp_subscripts[i - 1] != c) &&
+            (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c) &&
+            !out_ellipse_set.test(c)) {
+          output_subscript.append(1, c);
+        }
+      }
+      subscripts += "->" + out_ellipse + output_subscript;
+    }
+  }
+
+  // Build output string if does not exist
+  std::vector<std::string> ret(2);
+  if (subscripts.find("->") != std::string::npos) {
+    ret = split(subscripts, "->");
+  } else {
+    ret[0] = subscripts;
+    ret[1] = "";
+    // Build output subscripts
+    std::string tmp_subscripts = subscripts;
+    size_t len_tmp_subscripts = tmp_subscripts.length();
+    std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
+    for (size_t i = 0; i < len_tmp_subscripts; ++i) {
+      const char& c = tmp_subscripts[i];
+      if (c == ',') {
+        continue;
+      }
+      CHECK(einsum_symbols_set.test(c))
+        << "Character " << c
+        << " is not a valid symbol.";
+      if ((i == 0 || tmp_subscripts[i - 1] != c) &&
+          (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c)) {
+        ret[1].append(1, c);
+      }
+    }
+  }
+
+  // Make sure output subscripts are in the input
+  std::bitset<MAXAXIS> input_subscripts_set = str2set(ret[0]);
+  for (const char& c : ret[1]) {
+    CHECK(input_subscripts_set.test(c))
+      << "Output character " << c
+      << " did not appear in the input";
+  }
+
+  // Make sure number operands is equivalent to the number of terms
+  CHECK_EQ(std::count(ret[0].begin(), ret[0].end(), ',') + 1, operands.size())
+    << "Number of einsum subscripts must be equal to the "
+    << "number of operands.";
+
+  return ret;
+}
+
+inline Array<PrimExpr> NumpyEinsumShape(std::string subscripts, std::vector<Array<PrimExpr>>& operands) {
+  // Parsing
+  std::vector<std::string> parsed_subscripts = _parse_einsum_input(subscripts, operands);
+
+  // Build a few useful list and sets
+  std::vector<std::string> input_list = split(parsed_subscripts[0], ",");
+  size_t isize = input_list.size();
+
+  // Get length of each unique dimension and ensure all dimensions are correct
+  int dimension_dict[MAXAXIS];
+  memset(dimension_dict, -1, sizeof(dimension_dict));
+  for (size_t i = 0; i < isize; ++i) {
+    const std::string& term = input_list[i];
+    const Array<PrimExpr>& sh = operands[i];
+    CHECK_EQ(sh.size(), term.length())
+      << "Einstein sum subscript " << input_list[i]
+      << " does not contain the "
+      << "correct number of indices for operand " << i << ".";
+    size_t len_term = term.length();
+    for (size_t j = 0; j < len_term; ++j) {
+      int64_t dim = GetConstInt(sh[j]);
+      const char& c = term[j];
+
+      if (dimension_dict[static_cast<int>(c)] != -1) {
+        // For broadcasting cases we always want the largest dim size
+        if (dimension_dict[static_cast<int>(c)] == 1) {
+          dimension_dict[static_cast<int>(c)] = dim;
+        }
+        CHECK(dim == 1 || dim == dimension_dict[static_cast<int>(c)])
+          << "Size of label '" << c
+          << "' for operand  " << i
+          << " (" << dimension_dict[static_cast<int>(c)]
+          << ") does not match previous terms ("
+          << dim << ").";
+      } else {
+        dimension_dict[static_cast<int>(c)] = dim;
+      }
+    }
+  }
+
+  // Get oshape
+  const std::string& output_str = parsed_subscripts[1];
+  size_t odim = output_str.size();
+  Array<PrimExpr> oshape(odim, -1);
+  for (size_t i = 0; i < odim; ++i) {
+    oshape.Set(i, dimension_dict[static_cast<int>(output_str[i])]);
+  }
+  // Neglecting oshape assign check temporally
+  return oshape;
+}
+
+// inline Tensor einsum(const Array<Tensor> inputs, const std::string& subscripts_str, int optimize = 0,
+//                      std::string name = "T_einsum", std::string tag = kMatMul) {
+inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inputs,
+                     std::string name = "T_einsum", std::string tag = kMatMul) {
+  // able to compute op: trace, diag, sum, transpose, matmul, dot, inner, outer, multiply, tensordot
+  bool back = false;
+  const char* subscripts = subscripts_str.data();
+  const int nop = inputs.size();
+
+  /* Step 1: Parse the subscripts string into label_counts and op_labels */
+  int iop, idim, min_label = 127, max_label = 0;
+  char label_counts[128], op_labels[NPY_MAXARGS][NPY_MAXDIMS];
+  memset(label_counts, 0, sizeof(label_counts));
+  for (iop = 0; iop < nop; ++iop) {
+    int length = static_cast<int>(strcspn(subscripts, ",-"));
+
+    CHECK(!(iop == nop - 1 && subscripts[length] == ','))
+      << "more operands provided to einstein sum function "
+      << "than specified in the subscripts string";
+    CHECK(!(iop < nop-1 && subscripts[length] != ','))
+      << "fewer operands provided to einstein sum function "
+      << "than specified in the subscripts string";
+    parse_operand_subscripts(subscripts, length,
+                                      inputs[iop + back].ndim(),
+                                      iop, op_labels[iop], label_counts,
+                                      &min_label, &max_label);
+
+    /* Move subscripts to the start of the labels for the next op */
+    subscripts += length;
+    if (iop < nop - 1) {
+      subscripts++;
+    }
+  }
+  /*
+   * Find the number of broadcast dimensions, which is the maximum
+   * number of labels == 0 in an op_labels array.
+   */
+  int ndim_broadcast = 0;
+  for (iop = 0; iop < nop; ++iop) {
+    int count_zeros = 0;
+    int ndim;
+    char *labels = op_labels[iop];
+
+    ndim = inputs[iop + back].ndim();
+    for (idim = 0; idim < ndim; ++idim) {
+      if (labels[idim] == 0) {
+        ++count_zeros;
+      }
+    }
+
+    if (count_zeros > ndim_broadcast) {
+      ndim_broadcast = count_zeros;
+    }
+  }
+
+  /*
+   * If there is no output signature, fill output_labels and ndim_output
+   * using each label that appeared once, in alphabetical order.
+   */
+  int label, ndim_output;
+  char output_labels[NPY_MAXDIMS];
+  if (subscripts[0] == '\0') {
+    /* If no output was specified, always broadcast left, as usual. */
+    for (ndim_output = 0; ndim_output < ndim_broadcast; ++ndim_output) {
+      output_labels[ndim_output] = 0;
+    }
+    for (label = min_label; label <= max_label; ++label) {
+      if (label_counts[label] == 1) {
+        CHECK(ndim_output < NPY_MAXDIMS)
+          << "einstein sum subscript string has too many "
+          << "distinct labels";
+        output_labels[ndim_output++] = label;
+      }
+    }
+  } else {
+    CHECK(subscripts[0] == '-' && subscripts[1] == '>')
+      << "einstein sum subscript string does not "
+      << "contain proper '->' output specified";
+    subscripts += 2;
+
+    /* Parse the output subscript string. */
+    ndim_output = parse_output_subscripts(subscripts, strlen(subscripts),
+                                          ndim_broadcast, label_counts,
+                                          output_labels);
+    CHECK_GE(ndim_output, 0);
+  }
+
+  /*
+   * Step 2:
+   * Process all the input ops, combining dimensions into their
+   * diagonal where specified.
+   */
+  std::vector<Array<PrimExpr>> opshape(nop), opstride_true(nop);
+  for (iop = 0; iop < nop; ++iop) {
+    char *labels = op_labels[iop];
+    int combine, ndim;
+
+    ndim = inputs[iop + back].ndim();
+
+    /*
+     * Check whether any dimensions need to be combined
+     *
+     * The char type may be either signed or unsigned, we
+     * need it to be signed here.
+     */
+    combine = 0;
+    for (idim = 0; idim < ndim; ++idim) {
+      if ((signed char)labels[idim] < 0) {
+        combine++;
+      }
+    }
+    /* If any dimensions are combined, create a view which combines them */
+    if (combine) {
+      Array<PrimExpr> tshape(static_cast<size_t>(ndim - combine), -1);
+      Array<PrimExpr> tstride(static_cast<size_t>(ndim - combine), -1);
+      get_combined_dims_view(inputs[iop + back], iop, labels,
+                             tshape, tstride);
+      opshape[iop] = tshape;
+      opstride_true[iop] = tstride;
+    } else {
+      /* No combining needed */
+      opshape[iop] = inputs[iop + back]->shape;
+      opstride_true[iop] = get_stride(opshape[iop]);
+    }
+  }
+  /*
+   * Step 3:
+   * Set up the labels for the iterator (output + combined labels).
+   * Can just share the output_labels memory, because iter_labels
+   * is output_labels with some more labels appended.
+   */
+  char *iter_labels = output_labels;
+  int ndim_iter = ndim_output;
+  for (label = min_label; label <= max_label; ++label) {
+    if (label_counts[label] > 0 &&
+        memchr(output_labels, label, ndim_output) == nullptr) {
+      CHECK(ndim_iter < NPY_MAXDIMS)
+        << "too many subscripts in einsum";
+      iter_labels[ndim_iter++] = label;
+    }
+  }
+  /* Step 4: Set up the op_axes for the iterator */
+  Array<PrimExpr> itershape(static_cast<size_t>(ndim_iter), -1);
+  std::vector<Array<PrimExpr>> iterstride(nop + 1, Array<PrimExpr>(static_cast<size_t>(ndim_iter), 0));
+
+  // output_shape
+  std::vector<Array<PrimExpr>> operands;
+  for(size_t i = 0; i < inputs.size(); i++) {
+    operands.push_back(inputs[i]->shape);
+  }
+  Array<PrimExpr> oshape = NumpyEinsumShape(subscripts_str, operands);
+  Array<PrimExpr> ostride_true = get_stride(oshape);
+  Array<PrimExpr> reduceshape;
+  std::vector<Array<PrimExpr>> remainshape(nop);
+  int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
+  int *op_axes[NPY_MAXARGS];
+  for (iop = 0; iop < nop; ++iop) {
+    op_axes[iop] = op_axes_arrays[iop];
+    CHECK_GE(prepare_op_axes(opshape[iop].size(), iop, op_labels[iop],
+             op_axes[iop], ndim_iter, iter_labels), 0);
+    for (idim = 0; idim < ndim_iter; idim++) {
+      if (op_axes[iop][idim] != -1) {
+        iterstride[iop].Set(idim, opstride_true[iop][op_axes[iop][idim]]);
+        if (GetConstInt(itershape[idim]) != -1) {
+          if (GetConstInt(itershape[idim]) == 1) {
+            itershape.Set(idim, opshape[iop][op_axes[iop][idim]]);
+          }
+        } else {
+          itershape.Set(idim, opshape[iop][op_axes[iop][idim]]);
+        }
+      }
+    }
+  }
+  for (idim = 0; idim < ndim_output; ++idim) {
+    iterstride[nop].Set(idim, ostride_true[idim]);
+  }
+  reduceshape = Array<PrimExpr>(static_cast<size_t>(ndim_iter - ndim_output), 0);
+  for (idim = ndim_output; idim < ndim_iter; ++idim) {
+    reduceshape.Set(idim - ndim_output, itershape[idim]);
+  }
+  for (iop = 0; iop < nop; iop++) {
+    Array<Integer> rsh;
+    for (idim = 0; idim < ndim_iter; idim++) {
+
+      if (op_axes_arrays[iop][idim] == -1) {
+        rsh.push_back(GetConstInt(itershape[idim]));
+      } else {
+        if (GetConstInt(itershape[idim] != opshape[iop][op_axes_arrays[iop][idim]])) {
+          rsh.push_back(GetConstInt(itershape[idim]));
+        }
+      }
+      
+    }
+    remainshape[iop] = Array<PrimExpr>(rsh.begin(), rsh.end());
+  }
+  // exclude the 0-dim case
+  if (ndim_iter == 0) {
+    ndim_iter = 1;
+  }
+  itershape = pad(itershape, ndim_iter);
+  for (iop = 0; iop <= nop; ++iop) {
+    iterstride[iop] = pad(iterstride[iop], ndim_iter);
+  }
+  // oshape = pad(oshape, ndim_iter);
+  reduceshape = pad(reduceshape, ndim_iter);
+  for (iop = 0; iop < nop; ++iop) {
+    opshape[iop] = pad(opshape[iop], ndim_iter);
+    remainshape[iop] = pad(remainshape[iop], ndim_iter);
+  }
+  // ostride and rstride
+  Array<Array<PrimExpr>> ostride;
+  Array<Array<PrimExpr>> rstride;
+
+  for (iop = 0; iop < nop; ++iop) {
+    Array<PrimExpr> otmp(static_cast<size_t>(ndim_iter), 0);
+    Array<PrimExpr> rtmp(static_cast<size_t>(ndim_iter), 0);
+    for (idim = 0; idim < ndim_iter; ++idim) {
+      otmp.Set(idim, idim < ndim_output ? iterstride[iop][idim] : 1);
+      rtmp.Set(idim, idim < ndim_iter - ndim_output ? iterstride[iop][idim+ndim_output] : 1);
+    }
+    ostride.push_back(otmp);
+    rstride.push_back(rtmp);
+  }
+
+
+  // func: input indices => return cooresponding value
+  auto func = [inputs, oshape, ostride, reduceshape, ndim_iter, 
+               rstride, nop](const Array<Var>& input_indices) -> PrimExpr {
+
+      // input_indices本身就是带纬度的索引
+    for( int rdim = 0; rdim < ndim_iter; ++rdim) {
+      if (GetConstInt(reduceshape[rdim]) == 0) {
+        return 0;  //
+      }
+    }
+    Array<PrimExpr> ridx = UnravelIndex(0, reduceshape);
+    // AType needs to fit the output data
+    PrimExpr sum = 0;
+    bool rec_flag = false;
+    do {
+      PrimExpr tmp = 1;
+      for (int iop = 0; iop < nop; ++iop) {
+        if (iop != -1) {
+            PrimExpr k = 0;
+            for(size_t i = 0; i < input_indices.size(); ++i) {
+              k += input_indices[i] * ostride[iop][i];
+            }
+            for(size_t i = 0; i < ridx.size(); ++i) {
+              k += ridx[i] * rstride[iop][i];
+            }
+           Array<PrimExpr> temp_indices = UnravelIndex(k, inputs[iop]->shape);
+           tmp = tmp * inputs[iop](temp_indices);
+        }
+      }
+      sum += tmp;
+      
+      ridx.Set(ridx.size() - 1, ridx[ridx.size()-1] + 1);
+      for (int i = static_cast<int>(ridx.size() - 1); (i > 0) && GetConstInt(ridx[i] >= reduceshape[i]); --i){
+        ridx.Set(i, ridx[i] - reduceshape[i]);
+        ridx.Set(i-1, ridx[i-1] + 1);
+      }
+      rec_flag = GetConstInt(ridx[0] < reduceshape[0]);
+    } while(rec_flag);
+    return sum;
+  };
+
+  return compute(oshape, func, name, tag);
+}
+
 inline Tensor arange(const PrimExpr& start, const PrimExpr& stop, const PrimExpr& step,
                      DataType dtype, std::string name = "T_arange", std::string tag = kInjective) {
   PrimExpr num_elem = tvm::cast(
diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 6ddbc73e4666..6c3661288e3f 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -421,6 +421,7 @@ def take(a, indices, axis=None, mode="clip"):
     -------
     ret : tvm.te.Tensor
     """
+    print("type is ", type(mode))
     if axis is None:
         return cpp.take(a, indices, mode)
     return cpp.take(a, indices, int(axis), mode)
@@ -538,6 +539,12 @@ def tensordot(a, b, axes):
     return cpp.tensordot(a, b, axes[0], axes[1])
 
 
+def einsum(subscripts, a_tuple):
+    # Grab non-einsum kwargs; do not optimize by default.
+    # Temporally not implement the optimized einsum
+    return cpp.einsum(subscripts, a_tuple)
+
+
 def arange(start, stop=None, step=1, dtype="float32"):
     """Creates a tensor with evenly spaced values within a given interval.
 
diff --git a/src/topi/transform.cc b/src/topi/transform.cc
index e1e3988f6400..29c246f6ea82 100644
--- a/src/topi/transform.cc
+++ b/src/topi/transform.cc
@@ -86,6 +86,8 @@ TVM_REGISTER_GLOBAL("topi.layout_transform").set_body([](TVMArgs args, TVMRetVal
 });
 
 TVM_REGISTER_GLOBAL("topi.take").set_body([](TVMArgs args, TVMRetValue* rv) {
+  std::printf("%d\n", args[2].type_code());
+
   if (args.size() == 3) {
     std::string mode = args[2];
     *rv = take(args[0], args[1], mode);
@@ -165,6 +167,12 @@ TVM_REGISTER_GLOBAL("topi.tensordot").set_body([](TVMArgs args, TVMRetValue* rv)
   }
 });
 
+TVM_REGISTER_GLOBAL("topi.einsum").set_body([](TVMArgs args, TVMRetValue* rv) {
+  std::printf("%d\n", args[0].type_code());
+  std::printf("%d\n", args[1].type_code());
+  *rv = einsum(args[0], args[1]);
+});
+
 TVM_REGISTER_GLOBAL("topi.strided_slice").set_body([](TVMArgs args, TVMRetValue* rv) {
   *rv = strided_slice(args[0], args[1], args[2], args[3], args[4]);
 });

From c61eb6e32a026a858c32cc80d2cf615734366cdb Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Mon, 31 Aug 2020 22:32:52 +0800
Subject: [PATCH 02/14] * test

---
 3q                           | 28 ----------------------------
 include/tvm/topi/transform.h | 17 ++++++++++++-----
 python/tvm/topi/transform.py |  1 -
 src/topi/transform.cc        |  4 ----
 4 files changed, 12 insertions(+), 38 deletions(-)
 delete mode 100644 3q

diff --git a/3q b/3q
deleted file mode 100644
index 57d1462d5996..000000000000
--- a/3q
+++ /dev/null
@@ -1,28 +0,0 @@
-pick 6a175f57a * debug
-s caf7f5e64 * d
-
-# Rebase 4b61a9764..caf7f5e64 onto 4b61a9764 (2 commands)
-#
-# Commands:
-# p, pick <commit> = use commit
-# r, reword <commit> = use commit, but edit the commit message
-# e, edit <commit> = use commit, but stop for amending
-# s, squash <commit> = use commit, but meld into previous commit
-# f, fixup <commit> = like "squash", but discard this commit's log message
-# x, exec <command> = run command (the rest of the line) using shell
-# b, break = stop here (continue rebase later with 'git rebase --continue')
-# d, drop <commit> = remove commit
-# l, label <label> = label current HEAD with a name
-# t, reset <label> = reset HEAD to a label
-# m, merge [-C <commit> | -c <commit>] <label> [# <oneline>]
-# .       create a merge commit using the original merge commit's
-# .       message (or the oneline, if no original merge commit was
-# .       specified). Use -c <commit> to reword the commit message.
-#
-# These lines can be re-ordered; they are executed from top to bottom.
-#
-# If you remove a line here THAT COMMIT WILL BE LOST.
-#
-# However, if you remove everything, the rebase will be aborted.
-#
-# Note that empty commits are commented out
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index 94f1c50f26ec..8e52567706f5 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -1270,8 +1270,8 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, int axes = 2,
   Array<IterVar> iter_vars;
   for (int i = 0; i < axes; ++i)
     iter_vars.push_back(reduce_axis(Range(0, B->shape[i]), "k" + std::to_string(i)));
-  auto func = [&A, &B, &iter_vars, axes](const Array<Var>& input_indices) {
 
+  auto func = [&A, &B, &iter_vars, axes](const Array<Var>& input_indices) {
     Array<PrimExpr> A_indices(input_indices.begin(),
                               input_indices.begin() + (A->shape.size() - axes));
     for (auto& v : iter_vars) A_indices.push_back(v);
@@ -1285,9 +1285,8 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, int axes = 2,
     // Some passes don't like reductions with empty axis, so avoid it here
     if (iter_vars.empty())
       return A(A_indices) * B(B_indices);
-    else {
+    else
       return sum(A(A_indices) * B(B_indices), iter_vars);
-    }
   };
 
   return compute(output_shape, func, name, tag);
@@ -1898,8 +1897,16 @@ inline Array<PrimExpr> NumpyEinsumShape(std::string subscripts, std::vector<Arra
   return oshape;
 }
 
-// inline Tensor einsum(const Array<Tensor> inputs, const std::string& subscripts_str, int optimize = 0,
-//                      std::string name = "T_einsum", std::string tag = kMatMul) {
+/*!
+ * \brief Evaluates the Einstein summation convention on the operands.
+ *
+ * \param subscripts Specifies the subscripts for summation as comma separated list of subscript labels.
+ * \param inputs Arrays for the operation
+ * \param name The name of the operation
+ * \param tag The tag to mark the operation
+ *
+ * \return The calculation based on the Einstein summation convention.
+ */
 inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inputs,
                      std::string name = "T_einsum", std::string tag = kMatMul) {
   // able to compute op: trace, diag, sum, transpose, matmul, dot, inner, outer, multiply, tensordot
diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 6c3661288e3f..35075189d0d7 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -421,7 +421,6 @@ def take(a, indices, axis=None, mode="clip"):
     -------
     ret : tvm.te.Tensor
     """
-    print("type is ", type(mode))
     if axis is None:
         return cpp.take(a, indices, mode)
     return cpp.take(a, indices, int(axis), mode)
diff --git a/src/topi/transform.cc b/src/topi/transform.cc
index 29c246f6ea82..af47326eba36 100644
--- a/src/topi/transform.cc
+++ b/src/topi/transform.cc
@@ -86,8 +86,6 @@ TVM_REGISTER_GLOBAL("topi.layout_transform").set_body([](TVMArgs args, TVMRetVal
 });
 
 TVM_REGISTER_GLOBAL("topi.take").set_body([](TVMArgs args, TVMRetValue* rv) {
-  std::printf("%d\n", args[2].type_code());
-
   if (args.size() == 3) {
     std::string mode = args[2];
     *rv = take(args[0], args[1], mode);
@@ -168,8 +166,6 @@ TVM_REGISTER_GLOBAL("topi.tensordot").set_body([](TVMArgs args, TVMRetValue* rv)
 });
 
 TVM_REGISTER_GLOBAL("topi.einsum").set_body([](TVMArgs args, TVMRetValue* rv) {
-  std::printf("%d\n", args[0].type_code());
-  std::printf("%d\n", args[1].type_code());
   *rv = einsum(args[0], args[1]);
 });
 

From e8c90b6be05e72bac5ebd24dfffd4b10cda5e453 Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 2 Sep 2020 11:31:29 +0800
Subject: [PATCH 03/14] * Fix lint

---
 include/tvm/topi/transform.h | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index 8e52567706f5..0dc83a9b0803 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -1572,10 +1572,8 @@ inline void get_combined_dims_view(const Tensor& op, int iop,
       newstride.Set(icombine, stride[idim]);
     } else {
       /* Update the combined axis dimensions and strides */
-      int i = icombinemap[idim + combineoffset]; 
-      CHECK(!(
-        (combineoffset < 0) && GetConstInt(newshape[i] != 0 && newshape[i] != shape[idim])
-        ))
+      int i = icombinemap[idim + combineoffset];
+      CHECK(!((combineoffset < 0) && GetConstInt(newshape[i] != 0 && newshape[i] != shape[idim])))
         << "dimensions in operand " << iop
         << " for collapsing index '" << label
         << "' don't match (" << GetConstInt(newshape[i])
@@ -1846,7 +1844,8 @@ inline std::vector<std::string> _parse_einsum_input(
   return ret;
 }
 
-inline Array<PrimExpr> NumpyEinsumShape(std::string subscripts, std::vector<Array<PrimExpr>>& operands) {
+inline Array<PrimExpr> NumpyEinsumShape(std::string subscripts, 
+                                        std::vector<Array<PrimExpr>>& operands) {
   // Parsing
   std::vector<std::string> parsed_subscripts = _parse_einsum_input(subscripts, operands);
 
@@ -2048,11 +2047,12 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   }
   /* Step 4: Set up the op_axes for the iterator */
   Array<PrimExpr> itershape(static_cast<size_t>(ndim_iter), -1);
-  std::vector<Array<PrimExpr>> iterstride(nop + 1, Array<PrimExpr>(static_cast<size_t>(ndim_iter), 0));
+  std::vector<Array<PrimExpr>> iterstride(nop + 1, 
+                                Array<PrimExpr>(static_cast<size_t>(ndim_iter), 0));
 
   // output_shape
   std::vector<Array<PrimExpr>> operands;
-  for(size_t i = 0; i < inputs.size(); i++) {
+  for (size_t i = 0; i < inputs.size(); i++) {
     operands.push_back(inputs[i]->shape);
   }
   Array<PrimExpr> oshape = NumpyEinsumShape(subscripts_str, operands);
@@ -2088,7 +2088,6 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   for (iop = 0; iop < nop; iop++) {
     Array<Integer> rsh;
     for (idim = 0; idim < ndim_iter; idim++) {
-
       if (op_axes_arrays[iop][idim] == -1) {
         rsh.push_back(GetConstInt(itershape[idim]));
       } else {
@@ -2096,7 +2095,6 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
           rsh.push_back(GetConstInt(itershape[idim]));
         }
       }
-      
     }
     remainshape[iop] = Array<PrimExpr>(rsh.begin(), rsh.end());
   }
@@ -2131,11 +2129,9 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
 
 
   // func: input indices => return cooresponding value
-  auto func = [inputs, oshape, ostride, reduceshape, ndim_iter, 
+  auto func = [inputs, oshape, ostride, reduceshape, ndim_iter,
                rstride, nop](const Array<Var>& input_indices) -> PrimExpr {
-
-      // input_indices本身就是带纬度的索引
-    for( int rdim = 0; rdim < ndim_iter; ++rdim) {
+    for (int rdim = 0; rdim < ndim_iter; ++rdim) {
       if (GetConstInt(reduceshape[rdim]) == 0) {
         return 0;  //
       }
@@ -2149,10 +2145,10 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
       for (int iop = 0; iop < nop; ++iop) {
         if (iop != -1) {
             PrimExpr k = 0;
-            for(size_t i = 0; i < input_indices.size(); ++i) {
+            for (size_t i = 0; i < input_indices.size(); ++i) {
               k += input_indices[i] * ostride[iop][i];
             }
-            for(size_t i = 0; i < ridx.size(); ++i) {
+            for (size_t i = 0; i < ridx.size(); ++i) {
               k += ridx[i] * rstride[iop][i];
             }
            Array<PrimExpr> temp_indices = UnravelIndex(k, inputs[iop]->shape);
@@ -2162,12 +2158,13 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
       sum += tmp;
       
       ridx.Set(ridx.size() - 1, ridx[ridx.size()-1] + 1);
-      for (int i = static_cast<int>(ridx.size() - 1); (i > 0) && GetConstInt(ridx[i] >= reduceshape[i]); --i){
+      for (int i = static_cast<int>(ridx.size() - 1); (i > 0) &&
+                                  GetConstInt(ridx[i] >= reduceshape[i]); --i) {
         ridx.Set(i, ridx[i] - reduceshape[i]);
         ridx.Set(i-1, ridx[i-1] + 1);
       }
       rec_flag = GetConstInt(ridx[0] < reduceshape[0]);
-    } while(rec_flag);
+    } while (rec_flag);
     return sum;
   };
 

From 96a6e4e9b0bfcdd07e94d11ee4380d620bd572ec Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 2 Sep 2020 11:33:09 +0800
Subject: [PATCH 04/14] * Remove useless define

---
 include/tvm/topi/transform.h | 43 +++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 23 deletions(-)

diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index 0dc83a9b0803..5a433542270d 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -23,9 +23,6 @@
  */
 #ifndef TVM_TOPI_TRANSFORM_H_
 #define TVM_TOPI_TRANSFORM_H_
-#define NPY_MAXDIMS 16
-#define NPY_MAXARGS 16
-#define MAXAXIS 128
 
 #include <tvm/te/operation.h>
 #include <tvm/tir/data_layout.h>
@@ -1487,7 +1484,7 @@ inline int parse_output_subscripts(const char *subscripts, int length,
         << "in an input";
 
       /* Check that there is room in out_labels for this label. */
-      CHECK(ndim < NPY_MAXDIMS)
+      CHECK(ndim < 16)
         << "einstein sum subscripts string contains "
         << "too many subscripts in the output";
 
@@ -1502,7 +1499,7 @@ inline int parse_output_subscripts(const char *subscripts, int length,
         << "an ellipsis ('...') in the output";
 
       /* Check there is room in out_labels for broadcast dims. */
-      CHECK(ndim + ndim_broadcast <= NPY_MAXDIMS)
+      CHECK(ndim + ndim_broadcast <= 16)
         << "einstein sum subscripts string contains "
         << "too many subscripts in the output";
 
@@ -1528,12 +1525,12 @@ inline int parse_output_subscripts(const char *subscripts, int length,
   return ndim;
 }
 
-inline void get_combined_dims_view(const Tensor& op, int iop,
+inline void GetCombinedDimsView(const Tensor& op, int iop,
                                    char *labels,
                                    Array<PrimExpr> newshape,
                                    Array<PrimExpr> newstride) {
   int idim, ndim, icombine, combineoffset;
-  int icombinemap[NPY_MAXDIMS];
+  int icombinemap[16];
   int newdim;
 
   Array<PrimExpr> shape = op->shape;
@@ -1641,8 +1638,8 @@ inline int _count_substring(const std::string& str,
   return count;
 }
 
-inline std::bitset<MAXAXIS> str2set(const std::string& str) {
-  std::bitset<MAXAXIS> ret;
+inline std::bitset<128> str2set(const std::string& str) {
+  std::bitset<128> ret;
   for (const char& c : str) {
     ret.set(static_cast<int>(c));
   }
@@ -1667,7 +1664,7 @@ inline std::vector<std::string> _parse_einsum_input(
   const std::vector<Array<PrimExpr>>& operands) {
   const std::string einsum_symbols =
     "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
-  std::bitset<MAXAXIS> einsum_symbols_set;
+  std::bitset<128> einsum_symbols_set;
   for (const char& c : einsum_symbols) {
     einsum_symbols_set.set(c);
   }
@@ -1708,7 +1705,7 @@ inline std::vector<std::string> _parse_einsum_input(
                                                        c == '>';}),
                used.end());
 
-    std::bitset<MAXAXIS> used_set = str2set(used);
+    std::bitset<128> used_set = str2set(used);
     std::string ellipse_inds = "";
     for (const char& c : einsum_symbols) {
       if (!used_set.test(static_cast<int>(c))) {
@@ -1780,7 +1777,7 @@ inline std::vector<std::string> _parse_einsum_input(
       subscripts += "->" + output_sub;
     } else {
       // Special care for outputless ellipses
-      std::bitset<MAXAXIS> out_ellipse_set = str2set(out_ellipse);
+      std::bitset<128> out_ellipse_set = str2set(out_ellipse);
       std::string tmp_subscripts = subscripts, output_subscript = "";
       size_t len_tmp_subscripts = tmp_subscripts.length();
       std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
@@ -1829,7 +1826,7 @@ inline std::vector<std::string> _parse_einsum_input(
   }
 
   // Make sure output subscripts are in the input
-  std::bitset<MAXAXIS> input_subscripts_set = str2set(ret[0]);
+  std::bitset<128> input_subscripts_set = str2set(ret[0]);
   for (const char& c : ret[1]) {
     CHECK(input_subscripts_set.test(c))
       << "Output character " << c
@@ -1844,8 +1841,8 @@ inline std::vector<std::string> _parse_einsum_input(
   return ret;
 }
 
-inline Array<PrimExpr> NumpyEinsumShape(std::string subscripts, 
-                                        std::vector<Array<PrimExpr>>& operands) {
+inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts, 
+                                        const std::vector<Array<PrimExpr>>& operands) {
   // Parsing
   std::vector<std::string> parsed_subscripts = _parse_einsum_input(subscripts, operands);
 
@@ -1854,7 +1851,7 @@ inline Array<PrimExpr> NumpyEinsumShape(std::string subscripts,
   size_t isize = input_list.size();
 
   // Get length of each unique dimension and ensure all dimensions are correct
-  int dimension_dict[MAXAXIS];
+  int dimension_dict[128];
   memset(dimension_dict, -1, sizeof(dimension_dict));
   for (size_t i = 0; i < isize; ++i) {
     const std::string& term = input_list[i];
@@ -1915,7 +1912,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
 
   /* Step 1: Parse the subscripts string into label_counts and op_labels */
   int iop, idim, min_label = 127, max_label = 0;
-  char label_counts[128], op_labels[NPY_MAXARGS][NPY_MAXDIMS];
+  char label_counts[128], op_labels[16][16];
   memset(label_counts, 0, sizeof(label_counts));
   for (iop = 0; iop < nop; ++iop) {
     int length = static_cast<int>(strcspn(subscripts, ",-"));
@@ -1964,7 +1961,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
    * using each label that appeared once, in alphabetical order.
    */
   int label, ndim_output;
-  char output_labels[NPY_MAXDIMS];
+  char output_labels[16];
   if (subscripts[0] == '\0') {
     /* If no output was specified, always broadcast left, as usual. */
     for (ndim_output = 0; ndim_output < ndim_broadcast; ++ndim_output) {
@@ -1972,7 +1969,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     }
     for (label = min_label; label <= max_label; ++label) {
       if (label_counts[label] == 1) {
-        CHECK(ndim_output < NPY_MAXDIMS)
+        CHECK(ndim_output < 16)
           << "einstein sum subscript string has too many "
           << "distinct labels";
         output_labels[ndim_output++] = label;
@@ -2019,7 +2016,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     if (combine) {
       Array<PrimExpr> tshape(static_cast<size_t>(ndim - combine), -1);
       Array<PrimExpr> tstride(static_cast<size_t>(ndim - combine), -1);
-      get_combined_dims_view(inputs[iop + back], iop, labels,
+      GetCombinedDimsView(inputs[iop + back], iop, labels,
                              tshape, tstride);
       opshape[iop] = tshape;
       opstride_true[iop] = tstride;
@@ -2040,7 +2037,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   for (label = min_label; label <= max_label; ++label) {
     if (label_counts[label] > 0 &&
         memchr(output_labels, label, ndim_output) == nullptr) {
-      CHECK(ndim_iter < NPY_MAXDIMS)
+      CHECK(ndim_iter < 16)
         << "too many subscripts in einsum";
       iter_labels[ndim_iter++] = label;
     }
@@ -2059,8 +2056,8 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   Array<PrimExpr> ostride_true = get_stride(oshape);
   Array<PrimExpr> reduceshape;
   std::vector<Array<PrimExpr>> remainshape(nop);
-  int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
-  int *op_axes[NPY_MAXARGS];
+  int op_axes_arrays[16][16];
+  int *op_axes[16];
   for (iop = 0; iop < nop; ++iop) {
     op_axes[iop] = op_axes_arrays[iop];
     CHECK_GE(prepare_op_axes(opshape[iop].size(), iop, op_labels[iop],

From 08dcdc0345f940e5826b7e3cdf0a90528a6dbb7e Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 2 Sep 2020 19:20:57 +0800
Subject: [PATCH 05/14] * Move to einsum header file

---
 include/tvm/topi/einsum.h    | 871 +++++++++++++++++++++++++++++++++++
 include/tvm/topi/transform.h | 823 ---------------------------------
 src/topi/transform.cc        |   1 +
 3 files changed, 872 insertions(+), 823 deletions(-)
 create mode 100644 include/tvm/topi/einsum.h

diff --git a/include/tvm/topi/einsum.h b/include/tvm/topi/einsum.h
new file mode 100644
index 000000000000..7d797b96fbbd
--- /dev/null
+++ b/include/tvm/topi/einsum.h
@@ -0,0 +1,871 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * \file topi/einsum.h
+ * \brief Einstein summation op
+ */
+#ifndef TVM_TOPI_EINSUM_H_
+#define TVM_TOPI_EINSUM_H_
+
+#include <tvm/te/operation.h>
+#include <tvm/tir/data_layout.h>
+#include <tvm/topi/detail/constant_utils.h>
+#include <tvm/topi/detail/ravel_unravel.h>
+#include <tvm/topi/detail/tensor_utils.h>
+#include <tvm/topi/tags.h>
+
+#include <algorithm>
+#include <iterator>
+#include <string>
+#include <unordered_set>
+#include <vector>
+#include <bitset>
+
+namespace tvm {
+namespace topi {
+
+using namespace tvm::te;
+using namespace topi::detail;
+
+inline Array<PrimExpr> get_stride(const Array<PrimExpr> shape) {
+  size_t ndim = shape.size();
+  int prod = 1;
+  Array<PrimExpr> stride = Array<PrimExpr>(ndim, -1);
+  for (int i = ndim - 1; i >= 0; i--) {
+    stride.Set(i, if_then_else(shape[i] > 1, prod, 0));
+    prod = prod * GetConstInt(shape[i]);
+  }
+  return stride;
+}
+
+inline Array<PrimExpr> pad(const Array<PrimExpr> shape, int odim) {
+  int ndim = shape.size();
+  CHECK_GE(odim, ndim);
+  Array<PrimExpr> ret(static_cast<size_t>(odim), 1);
+  for (int idim = 0; idim < ndim; ++idim) {
+    ret.Set(idim, shape[idim]);
+  }
+  return ret;
+}
+
+inline int parse_operand_subscripts(const char *subscripts, int length,
+                                    int ndim, int iop, char *op_labels,
+                                    char *label_counts, int *min_label, int *max_label) {
+  int i;
+  int idim = 0;
+  int ellipsis = -1;
+
+  /* Process all labels for this operand */
+  for (i = 0; i < length; ++i) {
+    int label = subscripts[i];
+
+    /* A proper label for an axis. */
+    if (label > 0 && isalpha(label)) {
+      /* Check we don't exceed the operator dimensions. */
+      CHECK(idim < ndim)
+        << "einstein sum subscripts string contains "
+        << "too many subscripts for operand "
+        << iop;
+
+      op_labels[idim++] = label;
+      if (label < *min_label) {
+        *min_label = label;
+      }
+      if (label > *max_label) {
+        *max_label = label;
+      }
+      label_counts[label]++;
+    } else if (label == '.') {
+      /* The beginning of the ellipsis. */
+      /* Check it's a proper ellipsis. */
+      CHECK(!(ellipsis != -1 || i + 2 >= length
+              || subscripts[++i] != '.' || subscripts[++i] != '.'))
+        << "einstein sum subscripts string contains a "
+        << "'.' that is not part of an ellipsis ('...') "
+        << "in operand "
+        << iop;
+
+      ellipsis = idim;
+    } else {
+        CHECK(label == ' ')
+          << "invalid subscript '" << static_cast<char>(label)
+          << "' in einstein sum "
+          << "subscripts string, subscripts must "
+          << "be letters";
+    }
+  }
+
+  /* No ellipsis found, labels must match dimensions exactly. */
+  if (ellipsis == -1) {
+    CHECK(idim == ndim)
+      << "operand has more dimensions than subscripts "
+      << "given in einstein sum, but no '...' ellipsis "
+      << "provided to broadcast the extra dimensions.";
+  } else if (idim < ndim) {
+    /* Ellipsis found, may have to add broadcast dimensions. */
+    /* Move labels after ellipsis to the end. */
+    for (i = 0; i < idim - ellipsis; ++i) {
+      op_labels[ndim - i - 1] = op_labels[idim - i - 1];
+    }
+    /* Set all broadcast dimensions to zero. */
+    for (i = 0; i < ndim - idim; ++i) {
+      op_labels[ellipsis + i] = 0;
+    }
+  }
+
+  /*
+   * Find any labels duplicated for this operand, and turn them
+   * into negative offsets to the axis to merge with.
+   *
+   * In C, the char type may be signed or unsigned, but with
+   * twos complement arithmetic the char is ok either way here, and
+   * later where it matters the char is cast to a signed char.
+   */
+  for (idim = 0; idim < ndim - 1; ++idim) {
+    int label = op_labels[idim];
+    /* If it is a proper label, find any duplicates of it. */
+    if (label > 0) {
+      /* Search for the next matching label. */
+      char *next = reinterpret_cast<char*>(memchr(op_labels + idim + 1, label, ndim - idim - 1));
+
+      while (next != nullptr) {
+        /* The offset from next to op_labels[idim] (negative). */
+        *next = static_cast<char>((op_labels + idim) - next);
+        /* Search for the next matching label. */
+        next = reinterpret_cast<char*>(memchr(next + 1, label, op_labels + ndim - 1 - next));
+      }
+    }
+  }
+  return 0;
+}
+
+inline int parse_output_subscripts(const char *subscripts, int length,
+                                   int ndim_broadcast,
+                                   const char *label_counts, char *out_labels) {
+  int i, bdim;
+  int ndim = 0;
+  int ellipsis = 0;
+
+  /* Process all the output labels. */
+  for (i = 0; i < length; ++i) {
+    int label = subscripts[i];
+
+    /* A proper label for an axis. */
+    if (label > 0 && isalpha(label)) {
+      /* Check that it doesn't occur again. */
+      CHECK(memchr(subscripts + i + 1, label, length - i - 1) == nullptr)
+        << "einstein sum subscripts string includes "
+        << "output subscript '" << static_cast<char>(label)
+        << "' multiple times";
+
+      /* Check that it was used in the inputs. */
+      CHECK(label_counts[label] != 0)
+        << "einstein sum subscripts string included "
+        << "output subscript '" << static_cast<char>(label)
+        << "' which never appeared "
+        << "in an input";
+
+      /* Check that there is room in out_labels for this label. */
+      CHECK(ndim < 16)
+        << "einstein sum subscripts string contains "
+        << "too many subscripts in the output";
+
+      out_labels[ndim++] = label;
+    } else if (label == '.') {
+      /* The beginning of the ellipsis. */
+      /* Check it is a proper ellipsis. */
+      CHECK(!(ellipsis || i + 2 >= length
+              || subscripts[++i] != '.' || subscripts[++i] != '.'))
+        << "einstein sum subscripts string "
+        << "contains a '.' that is not part of "
+        << "an ellipsis ('...') in the output";
+
+      /* Check there is room in out_labels for broadcast dims. */
+      CHECK(ndim + ndim_broadcast <= 16)
+        << "einstein sum subscripts string contains "
+        << "too many subscripts in the output";
+
+      ellipsis = 1;
+      for (bdim = 0; bdim < ndim_broadcast; ++bdim) {
+        out_labels[ndim++] = 0;
+      }
+    } else {
+      CHECK(label == ' ')
+        << "invalid subscript '" << static_cast<char>(label)
+        << "' in einstein sum "
+        << "subscripts string, subscripts must "
+        << "be letters";
+    }
+  }
+
+  /* If no ellipsis was found there should be no broadcast dimensions. */
+  CHECK(!(!ellipsis && ndim_broadcast > 0))
+    << "output has more dimensions than subscripts "
+    << "given in einstein sum, but no '...' ellipsis "
+    << "provided to broadcast the extra dimensions.";
+
+  return ndim;
+}
+
+inline void GetCombinedDimsView(const Tensor& op, int iop,
+                                   char *labels,
+                                   Array<PrimExpr> newshape,
+                                   Array<PrimExpr> newstride) {
+  int idim, ndim, icombine, combineoffset;
+  int icombinemap[16];
+  int newdim;
+
+  Array<PrimExpr> shape = op->shape;
+  Array<PrimExpr> stride = get_stride(shape);
+  ndim = op.ndim();
+  newdim = newshape.size();
+
+  /* Initialize the dimensions and strides to zero */
+  for (idim = 0; idim < newdim; ++idim) {
+    newshape.Set(idim, 0);
+    newstride.Set(idim, 0);
+  }
+
+  /* Copy the dimensions and strides, except when collapsing */
+  icombine = 0;
+  for (idim = 0; idim < ndim; ++idim) {
+    /*
+     * The char type may be either signed or unsigned, we
+     * need it to be signed here.
+     */
+    int label = (signed char)labels[idim];
+    /* If this label says to merge axes, get the actual label */
+    if (label < 0) {
+      combineoffset = label;
+      label = labels[idim+label];
+    } else {
+      combineoffset = 0;
+      if (icombine != idim) {
+        labels[icombine] = labels[idim];
+      }
+      icombinemap[idim] = icombine;
+    }
+    /* If the label is 0, it's an unlabeled broadcast dimension */
+    if (label == 0) {
+      newshape.Set(icombine, shape[idim]);
+      newstride.Set(icombine, stride[idim]);
+    } else {
+      /* Update the combined axis dimensions and strides */
+      int i = icombinemap[idim + combineoffset];
+      CHECK(!((combineoffset < 0) && GetConstInt(newshape[i] != 0 && newshape[i] != shape[idim])))
+        << "dimensions in operand " << iop
+        << " for collapsing index '" << label
+        << "' don't match (" << GetConstInt(newshape[i])
+        << " != " << shape[idim] << ")";
+      newshape.Set(i, shape[idim]);
+      newstride.Set(i, newstride[i] + stride[idim]);
+    }
+
+    /* If the label didn't say to combine axes, increment dest i */
+    if (combineoffset == 0) {
+      icombine++;
+    }
+  }
+}
+
+inline static int prepare_op_axes(int ndim, int iop, char *labels,
+                                  int *axes, int ndim_iter, char *iter_labels) {
+  int i, label, ibroadcast;
+
+  ibroadcast = ndim-1;
+  for (i = ndim_iter-1; i >= 0; --i) {
+    label = iter_labels[i];
+    /*
+     * If it's an unlabeled broadcast dimension, choose
+     * the next broadcast dimension from the operand.
+     */
+    if (label == 0) {
+      while (ibroadcast >= 0 && labels[ibroadcast] != 0) {
+        --ibroadcast;
+      }
+      /*
+       * If we used up all the operand broadcast dimensions,
+       * extend it with a "newaxis"
+       */
+      if (ibroadcast < 0) {
+        axes[i] = -1;
+      } else {
+        /* Otherwise map to the broadcast axis */
+        axes[i] = ibroadcast;
+        --ibroadcast;
+      }
+    } else {
+      /* It's a labeled dimension, find the matching one */
+      char *match = reinterpret_cast<char*>(memchr(labels, label, ndim));
+      /* If the op doesn't have the label, broadcast it */
+      if (match == nullptr) {
+        axes[i] = -1;
+      } else {
+        /* Otherwise use it */
+        axes[i] = match - labels;
+      }
+    }
+  }
+  return 0;
+}
+
+inline int _count_substring(const std::string& str,
+                            const std::string& sub) {
+  int count = 0;
+  std::string::size_type pos = 0;
+  while ((pos = str.find(sub, pos)) != std::string::npos) {
+    ++count;
+    pos += sub.length();
+  }
+  return count;
+}
+
+inline std::bitset<128> str2set(const std::string& str) {
+  std::bitset<128> ret;
+  for (const char& c : str) {
+    ret.set(static_cast<int>(c));
+  }
+  return ret;
+}
+
+inline std::vector<std::string> split(const std::string& str,
+                                      const std::string& sub) {
+  std::string::size_type pos = 0;
+  std::string::size_type start = 0;
+  std::vector<std::string> ret;
+  while ((pos = str.find(sub, start)) != std::string::npos) {
+    ret.push_back(str.substr(start, pos - start));
+    start = pos + sub.length();
+  }
+  ret.push_back(str.substr(start));
+  return ret;
+}
+
+inline std::vector<std::string> _parse_einsum_input(
+  std::string subscripts,
+  const std::vector<Array<PrimExpr>>& operands) {
+  const std::string einsum_symbols =
+    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
+  std::bitset<128> einsum_symbols_set;
+  for (const char& c : einsum_symbols) {
+    einsum_symbols_set.set(c);
+  }
+
+  CHECK_NE(operands.size(), 0U)
+    << "No input operands";
+
+  auto end_pos = std::remove(subscripts.begin(), subscripts.end(), ' ');
+  subscripts.erase(end_pos, subscripts.end());
+
+  // Ensure all characters are valid
+  for (const char& c : subscripts) {
+    if (c == '.' || c == ',' || c == '-' || c == '>') {
+      continue;
+    }
+    CHECK(einsum_symbols_set.test(c))
+      << "Character " << c
+      << " is not a valid symbol.";
+  }
+
+  // Check for proper "->"
+  if (subscripts.find('-') != std::string::npos ||
+      subscripts.find('>') != std::string::npos) {
+    bool invalid = (std::count(subscripts.begin(), subscripts.end(), '-') > 1 ||
+                    std::count(subscripts.begin(), subscripts.end(), '>') > 1);
+    CHECK(!invalid && _count_substring(subscripts, "->") == 1)
+      << "Subscripts can only contain one '->'.";
+  }
+
+  // Parse ellipses
+  if (subscripts.find('.') != std::string::npos) {
+    std::string used = subscripts;
+    used.erase(std::remove_if(used.begin(),
+                              used.end(),
+                              [](const char& c){return c == '.' ||
+                                                       c == ',' ||
+                                                       c == '-' ||
+                                                       c == '>';}),
+               used.end());
+
+    std::bitset<128> used_set = str2set(used);
+    std::string ellipse_inds = "";
+    for (const char& c : einsum_symbols) {
+      if (!used_set.test(static_cast<int>(c))) {
+        ellipse_inds.append(1, c);
+      }
+    }
+    int longest = 0;
+    std::string input_tmp, output_sub;
+    std::vector<std::string> split_subscripts;
+    bool out_sub;
+
+    if (subscripts.find("->") != std::string::npos) {
+      std::vector<std::string> tmp = split(subscripts, "->");
+      input_tmp = tmp[0];
+      output_sub = tmp[1];
+      split_subscripts = split(input_tmp, ",");
+      out_sub = true;
+    } else {
+      split_subscripts = split(subscripts, ",");
+      out_sub = false;
+    }
+
+    size_t size_split_subscripts = split_subscripts.size();
+    subscripts = "";
+    for (size_t i = 0; i < size_split_subscripts; ++i) {
+      const std::string& sub = split_subscripts[i];
+      if (sub.find('.') != std::string::npos) {
+        CHECK_EQ(std::count(sub.begin(), sub.end(), '.'), 3)
+          << "Invalid Ellipses";
+        CHECK_EQ(_count_substring(sub, "..."), 1)
+          << "Invalid Ellipses";
+
+        // Take into account numerical values
+        int ellipse_count = 0;
+        if (operands[i].size() == 0) {
+          ellipse_count = 0;
+        } else {
+          ellipse_count = std::max(operands[i].size(), static_cast<size_t>(1));
+          ellipse_count -= sub.length() - 3;
+        }
+
+        if (ellipse_count > longest) {
+          longest = ellipse_count;
+        }
+
+        CHECK_GE(ellipse_count, 0)
+          << "Ellipses lengths do not match.";
+        if (ellipse_count == 0) {
+          split_subscripts[i].erase(sub.find("..."), 3);
+        } else {
+          std::string rep_inds = ellipse_inds.substr(ellipse_inds.length() - ellipse_count);
+          split_subscripts[i].replace(sub.find("..."), 3, rep_inds);
+        }
+      }
+      subscripts += split_subscripts[i];
+      if (i + 1 < size_split_subscripts) {
+        subscripts += ",";
+      }
+    }
+    std::string out_ellipse;
+    if (longest == 0) {
+      out_ellipse = "";
+    } else {
+      out_ellipse = ellipse_inds.substr(ellipse_inds.length() - longest);
+    }
+
+    if (out_sub) {
+      output_sub.replace(output_sub.find("..."), 3, out_ellipse);
+      subscripts += "->" + output_sub;
+    } else {
+      // Special care for outputless ellipses
+      std::bitset<128> out_ellipse_set = str2set(out_ellipse);
+      std::string tmp_subscripts = subscripts, output_subscript = "";
+      size_t len_tmp_subscripts = tmp_subscripts.length();
+      std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
+      for (size_t i = 0; i < len_tmp_subscripts; ++i) {
+        const char& c = tmp_subscripts[i];
+        if (c == ',') {
+          continue;
+        }
+        CHECK(einsum_symbols_set.test(c))
+          << "Character " << c
+          << " is not a valid symbol.";
+        if ((i == 0 || tmp_subscripts[i - 1] != c) &&
+            (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c) &&
+            !out_ellipse_set.test(c)) {
+          output_subscript.append(1, c);
+        }
+      }
+      subscripts += "->" + out_ellipse + output_subscript;
+    }
+  }
+
+  // Build output string if does not exist
+  std::vector<std::string> ret(2);
+  if (subscripts.find("->") != std::string::npos) {
+    ret = split(subscripts, "->");
+  } else {
+    ret[0] = subscripts;
+    ret[1] = "";
+    // Build output subscripts
+    std::string tmp_subscripts = subscripts;
+    size_t len_tmp_subscripts = tmp_subscripts.length();
+    std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
+    for (size_t i = 0; i < len_tmp_subscripts; ++i) {
+      const char& c = tmp_subscripts[i];
+      if (c == ',') {
+        continue;
+      }
+      CHECK(einsum_symbols_set.test(c))
+        << "Character " << c
+        << " is not a valid symbol.";
+      if ((i == 0 || tmp_subscripts[i - 1] != c) &&
+          (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c)) {
+        ret[1].append(1, c);
+      }
+    }
+  }
+
+  // Make sure output subscripts are in the input
+  std::bitset<128> input_subscripts_set = str2set(ret[0]);
+  for (const char& c : ret[1]) {
+    CHECK(input_subscripts_set.test(c))
+      << "Output character " << c
+      << " did not appear in the input";
+  }
+
+  // Make sure number operands is equivalent to the number of terms
+  CHECK_EQ(std::count(ret[0].begin(), ret[0].end(), ',') + 1, operands.size())
+    << "Number of einsum subscripts must be equal to the "
+    << "number of operands.";
+
+  return ret;
+}
+
+inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
+                                        const std::vector<Array<PrimExpr>>& operands) {
+  // Parsing
+  std::vector<std::string> parsed_subscripts = _parse_einsum_input(subscripts, operands);
+
+  // Build a few useful list and sets
+  std::vector<std::string> input_list = split(parsed_subscripts[0], ",");
+  size_t isize = input_list.size();
+
+  // Get length of each unique dimension and ensure all dimensions are correct
+  int dimension_dict[128];
+  memset(dimension_dict, -1, sizeof(dimension_dict));
+  for (size_t i = 0; i < isize; ++i) {
+    const std::string& term = input_list[i];
+    const Array<PrimExpr>& sh = operands[i];
+    CHECK_EQ(sh.size(), term.length())
+      << "Einstein sum subscript " << input_list[i]
+      << " does not contain the "
+      << "correct number of indices for operand " << i << ".";
+    size_t len_term = term.length();
+    for (size_t j = 0; j < len_term; ++j) {
+      int64_t dim = GetConstInt(sh[j]);
+      const char& c = term[j];
+
+      if (dimension_dict[static_cast<int>(c)] != -1) {
+        // For broadcasting cases we always want the largest dim size
+        if (dimension_dict[static_cast<int>(c)] == 1) {
+          dimension_dict[static_cast<int>(c)] = dim;
+        }
+        CHECK(dim == 1 || dim == dimension_dict[static_cast<int>(c)])
+          << "Size of label '" << c
+          << "' for operand  " << i
+          << " (" << dimension_dict[static_cast<int>(c)]
+          << ") does not match previous terms ("
+          << dim << ").";
+      } else {
+        dimension_dict[static_cast<int>(c)] = dim;
+      }
+    }
+  }
+
+  // Get oshape
+  const std::string& output_str = parsed_subscripts[1];
+  size_t odim = output_str.size();
+  Array<PrimExpr> oshape(odim, -1);
+  for (size_t i = 0; i < odim; ++i) {
+    oshape.Set(i, dimension_dict[static_cast<int>(output_str[i])]);
+  }
+  // Neglecting oshape assign check temporally
+  return oshape;
+}
+
+/*!
+ * \brief Evaluates the Einstein summation convention on the operands.
+ *
+ * \param subscripts Specifies the subscripts for summation as comma separated list of subscript labels.
+ * \param inputs Arrays for the operation
+ * \param name The name of the operation
+ * \param tag The tag to mark the operation
+ *
+ * \return The calculation based on the Einstein summation convention.
+ */
+inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inputs,
+                     std::string name = "T_einsum", std::string tag = kMatMul) {
+  // able to compute op: trace, diag, sum, transpose, matmul, dot, inner, outer, multiply, tensordot
+  bool back = false;
+  const char* subscripts = subscripts_str.data();
+  const int nop = inputs.size();
+
+  /* Step 1: Parse the subscripts string into label_counts and op_labels */
+  int iop, idim, min_label = 127, max_label = 0;
+  char label_counts[128], op_labels[16][16];
+  memset(label_counts, 0, sizeof(label_counts));
+  for (iop = 0; iop < nop; ++iop) {
+    int length = static_cast<int>(strcspn(subscripts, ",-"));
+
+    CHECK(!(iop == nop - 1 && subscripts[length] == ','))
+      << "more operands provided to einstein sum function "
+      << "than specified in the subscripts string";
+    CHECK(!(iop < nop-1 && subscripts[length] != ','))
+      << "fewer operands provided to einstein sum function "
+      << "than specified in the subscripts string";
+    parse_operand_subscripts(subscripts, length,
+                                      inputs[iop + back].ndim(),
+                                      iop, op_labels[iop], label_counts,
+                                      &min_label, &max_label);
+
+    /* Move subscripts to the start of the labels for the next op */
+    subscripts += length;
+    if (iop < nop - 1) {
+      subscripts++;
+    }
+  }
+  /*
+   * Find the number of broadcast dimensions, which is the maximum
+   * number of labels == 0 in an op_labels array.
+   */
+  int ndim_broadcast = 0;
+  for (iop = 0; iop < nop; ++iop) {
+    int count_zeros = 0;
+    int ndim;
+    char *labels = op_labels[iop];
+
+    ndim = inputs[iop + back].ndim();
+    for (idim = 0; idim < ndim; ++idim) {
+      if (labels[idim] == 0) {
+        ++count_zeros;
+      }
+    }
+
+    if (count_zeros > ndim_broadcast) {
+      ndim_broadcast = count_zeros;
+    }
+  }
+
+  /*
+   * If there is no output signature, fill output_labels and ndim_output
+   * using each label that appeared once, in alphabetical order.
+   */
+  int label, ndim_output;
+  char output_labels[16];
+  if (subscripts[0] == '\0') {
+    /* If no output was specified, always broadcast left, as usual. */
+    for (ndim_output = 0; ndim_output < ndim_broadcast; ++ndim_output) {
+      output_labels[ndim_output] = 0;
+    }
+    for (label = min_label; label <= max_label; ++label) {
+      if (label_counts[label] == 1) {
+        CHECK(ndim_output < 16)
+          << "einstein sum subscript string has too many "
+          << "distinct labels";
+        output_labels[ndim_output++] = label;
+      }
+    }
+  } else {
+    CHECK(subscripts[0] == '-' && subscripts[1] == '>')
+      << "einstein sum subscript string does not "
+      << "contain proper '->' output specified";
+    subscripts += 2;
+
+    /* Parse the output subscript string. */
+    ndim_output = parse_output_subscripts(subscripts, strlen(subscripts),
+                                          ndim_broadcast, label_counts,
+                                          output_labels);
+    CHECK_GE(ndim_output, 0);
+  }
+
+  /*
+   * Step 2:
+   * Process all the input ops, combining dimensions into their
+   * diagonal where specified.
+   */
+  std::vector<Array<PrimExpr>> opshape(nop), opstride_true(nop);
+  for (iop = 0; iop < nop; ++iop) {
+    char *labels = op_labels[iop];
+    int combine, ndim;
+
+    ndim = inputs[iop + back].ndim();
+
+    /*
+     * Check whether any dimensions need to be combined
+     *
+     * The char type may be either signed or unsigned, we
+     * need it to be signed here.
+     */
+    combine = 0;
+    for (idim = 0; idim < ndim; ++idim) {
+      if ((signed char)labels[idim] < 0) {
+        combine++;
+      }
+    }
+    /* If any dimensions are combined, create a view which combines them */
+    if (combine) {
+      Array<PrimExpr> tshape(static_cast<size_t>(ndim - combine), -1);
+      Array<PrimExpr> tstride(static_cast<size_t>(ndim - combine), -1);
+      GetCombinedDimsView(inputs[iop + back], iop, labels,
+                             tshape, tstride);
+      opshape[iop] = tshape;
+      opstride_true[iop] = tstride;
+    } else {
+      /* No combining needed */
+      opshape[iop] = inputs[iop + back]->shape;
+      opstride_true[iop] = get_stride(opshape[iop]);
+    }
+  }
+  /*
+   * Step 3:
+   * Set up the labels for the iterator (output + combined labels).
+   * Can just share the output_labels memory, because iter_labels
+   * is output_labels with some more labels appended.
+   */
+  char *iter_labels = output_labels;
+  int ndim_iter = ndim_output;
+  for (label = min_label; label <= max_label; ++label) {
+    if (label_counts[label] > 0 &&
+        memchr(output_labels, label, ndim_output) == nullptr) {
+      CHECK(ndim_iter < 16)
+        << "too many subscripts in einsum";
+      iter_labels[ndim_iter++] = label;
+    }
+  }
+  /* Step 4: Set up the op_axes for the iterator */
+  Array<PrimExpr> itershape(static_cast<size_t>(ndim_iter), -1);
+  std::vector<Array<PrimExpr>> iterstride(nop + 1,
+                                Array<PrimExpr>(static_cast<size_t>(ndim_iter), 0));
+
+  // output_shape
+  std::vector<Array<PrimExpr>> operands;
+  for (size_t i = 0; i < inputs.size(); i++) {
+    operands.push_back(inputs[i]->shape);
+  }
+  Array<PrimExpr> oshape = NumpyEinsumShape(subscripts_str, operands);
+  Array<PrimExpr> ostride_true = get_stride(oshape);
+  Array<PrimExpr> reduceshape;
+  std::vector<Array<PrimExpr>> remainshape(nop);
+  int op_axes_arrays[16][16];
+  int *op_axes[16];
+  for (iop = 0; iop < nop; ++iop) {
+    op_axes[iop] = op_axes_arrays[iop];
+    CHECK_GE(prepare_op_axes(opshape[iop].size(), iop, op_labels[iop],
+             op_axes[iop], ndim_iter, iter_labels), 0);
+    for (idim = 0; idim < ndim_iter; idim++) {
+      if (op_axes[iop][idim] != -1) {
+        iterstride[iop].Set(idim, opstride_true[iop][op_axes[iop][idim]]);
+        if (GetConstInt(itershape[idim]) != -1) {
+          if (GetConstInt(itershape[idim]) == 1) {
+            itershape.Set(idim, opshape[iop][op_axes[iop][idim]]);
+          }
+        } else {
+          itershape.Set(idim, opshape[iop][op_axes[iop][idim]]);
+        }
+      }
+    }
+  }
+  for (idim = 0; idim < ndim_output; ++idim) {
+    iterstride[nop].Set(idim, ostride_true[idim]);
+  }
+  reduceshape = Array<PrimExpr>(static_cast<size_t>(ndim_iter - ndim_output), 0);
+  for (idim = ndim_output; idim < ndim_iter; ++idim) {
+    reduceshape.Set(idim - ndim_output, itershape[idim]);
+  }
+  for (iop = 0; iop < nop; iop++) {
+    Array<Integer> rsh;
+    for (idim = 0; idim < ndim_iter; idim++) {
+      if (op_axes_arrays[iop][idim] == -1) {
+        rsh.push_back(GetConstInt(itershape[idim]));
+      } else {
+        if (GetConstInt(itershape[idim] != opshape[iop][op_axes_arrays[iop][idim]])) {
+          rsh.push_back(GetConstInt(itershape[idim]));
+        }
+      }
+    }
+    remainshape[iop] = Array<PrimExpr>(rsh.begin(), rsh.end());
+  }
+  // exclude the 0-dim case
+  if (ndim_iter == 0) {
+    ndim_iter = 1;
+  }
+  itershape = pad(itershape, ndim_iter);
+  for (iop = 0; iop <= nop; ++iop) {
+    iterstride[iop] = pad(iterstride[iop], ndim_iter);
+  }
+  // oshape = pad(oshape, ndim_iter);
+  reduceshape = pad(reduceshape, ndim_iter);
+  for (iop = 0; iop < nop; ++iop) {
+    opshape[iop] = pad(opshape[iop], ndim_iter);
+    remainshape[iop] = pad(remainshape[iop], ndim_iter);
+  }
+  // ostride and rstride
+  Array<Array<PrimExpr>> ostride;
+  Array<Array<PrimExpr>> rstride;
+
+  for (iop = 0; iop < nop; ++iop) {
+    Array<PrimExpr> otmp(static_cast<size_t>(ndim_iter), 0);
+    Array<PrimExpr> rtmp(static_cast<size_t>(ndim_iter), 0);
+    for (idim = 0; idim < ndim_iter; ++idim) {
+      otmp.Set(idim, idim < ndim_output ? iterstride[iop][idim] : 1);
+      rtmp.Set(idim, idim < ndim_iter - ndim_output ? iterstride[iop][idim+ndim_output] : 1);
+    }
+    ostride.push_back(otmp);
+    rstride.push_back(rtmp);
+  }
+
+
+  // func: input indices => return cooresponding value
+  auto func = [inputs, oshape, ostride, reduceshape, ndim_iter,
+               rstride, nop](const Array<Var>& input_indices) -> PrimExpr {
+    for (int rdim = 0; rdim < ndim_iter; ++rdim) {
+      if (GetConstInt(reduceshape[rdim]) == 0) {
+        return 0;  //
+      }
+    }
+    Array<PrimExpr> ridx = UnravelIndex(0, reduceshape);
+    // AType needs to fit the output data
+    PrimExpr sum = 0;
+    bool rec_flag = false;
+    do {
+      PrimExpr tmp = 1;
+      for (int iop = 0; iop < nop; ++iop) {
+        if (iop != -1) {
+            PrimExpr k = 0;
+            for (size_t i = 0; i < input_indices.size(); ++i) {
+              k += input_indices[i] * ostride[iop][i];
+            }
+            for (size_t i = 0; i < ridx.size(); ++i) {
+              k += ridx[i] * rstride[iop][i];
+            }
+           Array<PrimExpr> temp_indices = UnravelIndex(k, inputs[iop]->shape);
+           tmp = tmp * inputs[iop](temp_indices);
+        }
+      }
+      sum += tmp;
+      ridx.Set(ridx.size() - 1, ridx[ridx.size()-1] + 1);
+      for (int i = static_cast<int>(ridx.size() - 1); (i > 0) &&
+                                  GetConstInt(ridx[i] >= reduceshape[i]); --i) {
+        ridx.Set(i, ridx[i] - reduceshape[i]);
+        ridx.Set(i-1, ridx[i-1] + 1);
+      }
+      rec_flag = GetConstInt(ridx[0] < reduceshape[0]);
+    } while (rec_flag);
+    return sum;
+  };
+
+  return compute(oshape, func, name, tag);
+}
+
+}  // namespace topi
+}  // namespace tvm
+#endif  // TVM_TOPI_EINSUM_H_
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index 5a433542270d..0d7fc9d9f010 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -1345,829 +1345,6 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, Array<PrimExp
   return compute(output_shape, func, name, tag);
 }
 
-inline Array<PrimExpr> get_stride(const Array<PrimExpr> shape) {
-  size_t ndim = shape.size();
-  int prod = 1;
-  Array<PrimExpr> stride = Array<PrimExpr>(ndim, -1);
-  for (int i = ndim - 1; i >= 0; i--) {
-    stride.Set(i, if_then_else(shape[i] > 1, prod, 0));
-    prod = prod * GetConstInt(shape[i]);
-  }
-  return stride;
-}
-
-inline Array<PrimExpr> pad(const Array<PrimExpr> shape, int odim) {
-  int ndim = shape.size();
-  CHECK_GE(odim, ndim);
-  Array<PrimExpr> ret(static_cast<size_t>(odim), 1);
-  for (int idim = 0; idim < ndim; ++idim) {
-    ret.Set(idim, shape[idim]);
-  }
-  return ret;
-}
-
-inline int parse_operand_subscripts(const char *subscripts, int length,
-                                    int ndim, int iop, char *op_labels,
-                                    char *label_counts, int *min_label, int *max_label) {
-  int i;
-  int idim = 0;
-  int ellipsis = -1;
-
-  /* Process all labels for this operand */
-  for (i = 0; i < length; ++i) {
-    int label = subscripts[i];
-
-    /* A proper label for an axis. */
-    if (label > 0 && isalpha(label)) {
-      /* Check we don't exceed the operator dimensions. */
-      CHECK(idim < ndim)
-        << "einstein sum subscripts string contains "
-        << "too many subscripts for operand "
-        << iop;
-
-      op_labels[idim++] = label;
-      if (label < *min_label) {
-        *min_label = label;
-      }
-      if (label > *max_label) {
-        *max_label = label;
-      }
-      label_counts[label]++;
-    } else if (label == '.') {
-      /* The beginning of the ellipsis. */
-      /* Check it's a proper ellipsis. */
-      CHECK(!(ellipsis != -1 || i + 2 >= length
-              || subscripts[++i] != '.' || subscripts[++i] != '.'))
-        << "einstein sum subscripts string contains a "
-        << "'.' that is not part of an ellipsis ('...') "
-        << "in operand "
-        << iop;
-
-      ellipsis = idim;
-    } else {
-        CHECK(label == ' ')
-          << "invalid subscript '" << static_cast<char>(label)
-          << "' in einstein sum "
-          << "subscripts string, subscripts must "
-          << "be letters";
-    }
-  }
-
-  /* No ellipsis found, labels must match dimensions exactly. */
-  if (ellipsis == -1) {
-    CHECK(idim == ndim)
-      << "operand has more dimensions than subscripts "
-      << "given in einstein sum, but no '...' ellipsis "
-      << "provided to broadcast the extra dimensions.";
-  } else if (idim < ndim) {
-    /* Ellipsis found, may have to add broadcast dimensions. */
-    /* Move labels after ellipsis to the end. */
-    for (i = 0; i < idim - ellipsis; ++i) {
-      op_labels[ndim - i - 1] = op_labels[idim - i - 1];
-    }
-    /* Set all broadcast dimensions to zero. */
-    for (i = 0; i < ndim - idim; ++i) {
-      op_labels[ellipsis + i] = 0;
-    }
-  }
-
-  /*
-   * Find any labels duplicated for this operand, and turn them
-   * into negative offsets to the axis to merge with.
-   *
-   * In C, the char type may be signed or unsigned, but with
-   * twos complement arithmetic the char is ok either way here, and
-   * later where it matters the char is cast to a signed char.
-   */
-  for (idim = 0; idim < ndim - 1; ++idim) {
-    int label = op_labels[idim];
-    /* If it is a proper label, find any duplicates of it. */
-    if (label > 0) {
-      /* Search for the next matching label. */
-      char *next = reinterpret_cast<char*>(memchr(op_labels + idim + 1, label, ndim - idim - 1));
-
-      while (next != nullptr) {
-        /* The offset from next to op_labels[idim] (negative). */
-        *next = static_cast<char>((op_labels + idim) - next);
-        /* Search for the next matching label. */
-        next = reinterpret_cast<char*>(memchr(next + 1, label, op_labels + ndim - 1 - next));
-      }
-    }
-  }
-  return 0;
-}
-
-inline int parse_output_subscripts(const char *subscripts, int length,
-                                   int ndim_broadcast,
-                                   const char *label_counts, char *out_labels) {
-  int i, bdim;
-  int ndim = 0;
-  int ellipsis = 0;
-
-  /* Process all the output labels. */
-  for (i = 0; i < length; ++i) {
-    int label = subscripts[i];
-
-    /* A proper label for an axis. */
-    if (label > 0 && isalpha(label)) {
-      /* Check that it doesn't occur again. */
-      CHECK(memchr(subscripts + i + 1, label, length - i - 1) == nullptr)
-        << "einstein sum subscripts string includes "
-        << "output subscript '" << static_cast<char>(label)
-        << "' multiple times";
-
-      /* Check that it was used in the inputs. */
-      CHECK(label_counts[label] != 0)
-        << "einstein sum subscripts string included "
-        << "output subscript '" << static_cast<char>(label)
-        << "' which never appeared "
-        << "in an input";
-
-      /* Check that there is room in out_labels for this label. */
-      CHECK(ndim < 16)
-        << "einstein sum subscripts string contains "
-        << "too many subscripts in the output";
-
-      out_labels[ndim++] = label;
-    } else if (label == '.') {
-      /* The beginning of the ellipsis. */
-      /* Check it is a proper ellipsis. */
-      CHECK(!(ellipsis || i + 2 >= length
-              || subscripts[++i] != '.' || subscripts[++i] != '.'))
-        << "einstein sum subscripts string "
-        << "contains a '.' that is not part of "
-        << "an ellipsis ('...') in the output";
-
-      /* Check there is room in out_labels for broadcast dims. */
-      CHECK(ndim + ndim_broadcast <= 16)
-        << "einstein sum subscripts string contains "
-        << "too many subscripts in the output";
-
-      ellipsis = 1;
-      for (bdim = 0; bdim < ndim_broadcast; ++bdim) {
-        out_labels[ndim++] = 0;
-      }
-    } else {
-      CHECK(label == ' ')
-        << "invalid subscript '" << static_cast<char>(label)
-        << "' in einstein sum "
-        << "subscripts string, subscripts must "
-        << "be letters";
-    }
-  }
-
-  /* If no ellipsis was found there should be no broadcast dimensions. */
-  CHECK(!(!ellipsis && ndim_broadcast > 0))
-    << "output has more dimensions than subscripts "
-    << "given in einstein sum, but no '...' ellipsis "
-    << "provided to broadcast the extra dimensions.";
-
-  return ndim;
-}
-
-inline void GetCombinedDimsView(const Tensor& op, int iop,
-                                   char *labels,
-                                   Array<PrimExpr> newshape,
-                                   Array<PrimExpr> newstride) {
-  int idim, ndim, icombine, combineoffset;
-  int icombinemap[16];
-  int newdim;
-
-  Array<PrimExpr> shape = op->shape;
-  Array<PrimExpr> stride = get_stride(shape);
-  ndim = op.ndim();
-  newdim = newshape.size();
-
-  /* Initialize the dimensions and strides to zero */
-  for (idim = 0; idim < newdim; ++idim) {
-    newshape.Set(idim, 0);
-    newstride.Set(idim, 0);
-  }
-
-  /* Copy the dimensions and strides, except when collapsing */
-  icombine = 0;
-  for (idim = 0; idim < ndim; ++idim) {
-    /*
-     * The char type may be either signed or unsigned, we
-     * need it to be signed here.
-     */
-    int label = (signed char)labels[idim];
-    /* If this label says to merge axes, get the actual label */
-    if (label < 0) {
-      combineoffset = label;
-      label = labels[idim+label];
-    } else {
-      combineoffset = 0;
-      if (icombine != idim) {
-        labels[icombine] = labels[idim];
-      }
-      icombinemap[idim] = icombine;
-    }
-    /* If the label is 0, it's an unlabeled broadcast dimension */
-    if (label == 0) {
-      newshape.Set(icombine, shape[idim]);
-      newstride.Set(icombine, stride[idim]);
-    } else {
-      /* Update the combined axis dimensions and strides */
-      int i = icombinemap[idim + combineoffset];
-      CHECK(!((combineoffset < 0) && GetConstInt(newshape[i] != 0 && newshape[i] != shape[idim])))
-        << "dimensions in operand " << iop
-        << " for collapsing index '" << label
-        << "' don't match (" << GetConstInt(newshape[i])
-        << " != " << shape[idim] << ")";
-      newshape.Set(i, shape[idim]);
-      newstride.Set(i, newstride[i] + stride[idim]);
-    }
-
-    /* If the label didn't say to combine axes, increment dest i */
-    if (combineoffset == 0) {
-      icombine++;
-    }
-  }
-}
-
-inline static int prepare_op_axes(int ndim, int iop, char *labels,
-                                  int *axes, int ndim_iter, char *iter_labels) {
-  int i, label, ibroadcast;
-
-  ibroadcast = ndim-1;
-  for (i = ndim_iter-1; i >= 0; --i) {
-    label = iter_labels[i];
-    /*
-     * If it's an unlabeled broadcast dimension, choose
-     * the next broadcast dimension from the operand.
-     */
-    if (label == 0) {
-      while (ibroadcast >= 0 && labels[ibroadcast] != 0) {
-        --ibroadcast;
-      }
-      /*
-       * If we used up all the operand broadcast dimensions,
-       * extend it with a "newaxis"
-       */
-      if (ibroadcast < 0) {
-        axes[i] = -1;
-      } else {
-        /* Otherwise map to the broadcast axis */
-        axes[i] = ibroadcast;
-        --ibroadcast;
-      }
-    } else {
-      /* It's a labeled dimension, find the matching one */
-      char *match = reinterpret_cast<char*>(memchr(labels, label, ndim));
-      /* If the op doesn't have the label, broadcast it */
-      if (match == nullptr) {
-        axes[i] = -1;
-      } else {
-        /* Otherwise use it */
-        axes[i] = match - labels;
-      }
-    }
-  }
-  return 0;
-}
-
-inline int _count_substring(const std::string& str,
-                            const std::string& sub) {
-  int count = 0;
-  std::string::size_type pos = 0;
-  while ((pos = str.find(sub, pos)) != std::string::npos) {
-    ++count;
-    pos += sub.length();
-  }
-  return count;
-}
-
-inline std::bitset<128> str2set(const std::string& str) {
-  std::bitset<128> ret;
-  for (const char& c : str) {
-    ret.set(static_cast<int>(c));
-  }
-  return ret;
-}
-
-inline std::vector<std::string> split(const std::string& str,
-                                      const std::string& sub) {
-  std::string::size_type pos = 0;
-  std::string::size_type start = 0;
-  std::vector<std::string> ret;
-  while ((pos = str.find(sub, start)) != std::string::npos) {
-    ret.push_back(str.substr(start, pos - start));
-    start = pos + sub.length();
-  }
-  ret.push_back(str.substr(start));
-  return ret;
-}
-
-inline std::vector<std::string> _parse_einsum_input(
-  std::string subscripts,
-  const std::vector<Array<PrimExpr>>& operands) {
-  const std::string einsum_symbols =
-    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
-  std::bitset<128> einsum_symbols_set;
-  for (const char& c : einsum_symbols) {
-    einsum_symbols_set.set(c);
-  }
-
-  CHECK_NE(operands.size(), 0U)
-    << "No input operands";
-
-  auto end_pos = std::remove(subscripts.begin(), subscripts.end(), ' ');
-  subscripts.erase(end_pos, subscripts.end());
-
-  // Ensure all characters are valid
-  for (const char& c : subscripts) {
-    if (c == '.' || c == ',' || c == '-' || c == '>') {
-      continue;
-    }
-    CHECK(einsum_symbols_set.test(c))
-      << "Character " << c
-      << " is not a valid symbol.";
-  }
-
-  // Check for proper "->"
-  if (subscripts.find('-') != std::string::npos ||
-      subscripts.find('>') != std::string::npos) {
-    bool invalid = (std::count(subscripts.begin(), subscripts.end(), '-') > 1 ||
-                    std::count(subscripts.begin(), subscripts.end(), '>') > 1);
-    CHECK(!invalid && _count_substring(subscripts, "->") == 1)
-      << "Subscripts can only contain one '->'.";
-  }
-
-  // Parse ellipses
-  if (subscripts.find('.') != std::string::npos) {
-    std::string used = subscripts;
-    used.erase(std::remove_if(used.begin(),
-                              used.end(),
-                              [](const char& c){return c == '.' ||
-                                                       c == ',' ||
-                                                       c == '-' ||
-                                                       c == '>';}),
-               used.end());
-
-    std::bitset<128> used_set = str2set(used);
-    std::string ellipse_inds = "";
-    for (const char& c : einsum_symbols) {
-      if (!used_set.test(static_cast<int>(c))) {
-        ellipse_inds.append(1, c);
-      }
-    }
-    int longest = 0;
-    std::string input_tmp, output_sub;
-    std::vector<std::string> split_subscripts;
-    bool out_sub;
-
-    if (subscripts.find("->") != std::string::npos) {
-      std::vector<std::string> tmp = split(subscripts, "->");
-      input_tmp = tmp[0];
-      output_sub = tmp[1];
-      split_subscripts = split(input_tmp, ",");
-      out_sub = true;
-    } else {
-      split_subscripts = split(subscripts, ",");
-      out_sub = false;
-    }
-
-    size_t size_split_subscripts = split_subscripts.size();
-    subscripts = "";
-    for (size_t i = 0; i < size_split_subscripts; ++i) {
-      const std::string& sub = split_subscripts[i];
-      if (sub.find('.') != std::string::npos) {
-        CHECK_EQ(std::count(sub.begin(), sub.end(), '.'), 3)
-          << "Invalid Ellipses";
-        CHECK_EQ(_count_substring(sub, "..."), 1)
-          << "Invalid Ellipses";
-
-        // Take into account numerical values
-        int ellipse_count = 0;
-        if (operands[i].size() == 0) {
-          ellipse_count = 0;
-        } else {
-          ellipse_count = std::max(operands[i].size(), static_cast<size_t>(1));
-          ellipse_count -= sub.length() - 3;
-        }
-
-        if (ellipse_count > longest) {
-          longest = ellipse_count;
-        }
-
-        CHECK_GE(ellipse_count, 0)
-          << "Ellipses lengths do not match.";
-        if (ellipse_count == 0) {
-          split_subscripts[i].erase(sub.find("..."), 3);
-        } else {
-          std::string rep_inds = ellipse_inds.substr(ellipse_inds.length() - ellipse_count);
-          split_subscripts[i].replace(sub.find("..."), 3, rep_inds);
-        }
-      }
-      subscripts += split_subscripts[i];
-      if (i + 1 < size_split_subscripts) {
-        subscripts += ",";
-      }
-    }
-    std::string out_ellipse;
-    if (longest == 0) {
-      out_ellipse = "";
-    } else {
-      out_ellipse = ellipse_inds.substr(ellipse_inds.length() - longest);
-    }
-
-    if (out_sub) {
-      output_sub.replace(output_sub.find("..."), 3, out_ellipse);
-      subscripts += "->" + output_sub;
-    } else {
-      // Special care for outputless ellipses
-      std::bitset<128> out_ellipse_set = str2set(out_ellipse);
-      std::string tmp_subscripts = subscripts, output_subscript = "";
-      size_t len_tmp_subscripts = tmp_subscripts.length();
-      std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
-      for (size_t i = 0; i < len_tmp_subscripts; ++i) {
-        const char& c = tmp_subscripts[i];
-        if (c == ',') {
-          continue;
-        }
-        CHECK(einsum_symbols_set.test(c))
-          << "Character " << c
-          << " is not a valid symbol.";
-        if ((i == 0 || tmp_subscripts[i - 1] != c) &&
-            (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c) &&
-            !out_ellipse_set.test(c)) {
-          output_subscript.append(1, c);
-        }
-      }
-      subscripts += "->" + out_ellipse + output_subscript;
-    }
-  }
-
-  // Build output string if does not exist
-  std::vector<std::string> ret(2);
-  if (subscripts.find("->") != std::string::npos) {
-    ret = split(subscripts, "->");
-  } else {
-    ret[0] = subscripts;
-    ret[1] = "";
-    // Build output subscripts
-    std::string tmp_subscripts = subscripts;
-    size_t len_tmp_subscripts = tmp_subscripts.length();
-    std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
-    for (size_t i = 0; i < len_tmp_subscripts; ++i) {
-      const char& c = tmp_subscripts[i];
-      if (c == ',') {
-        continue;
-      }
-      CHECK(einsum_symbols_set.test(c))
-        << "Character " << c
-        << " is not a valid symbol.";
-      if ((i == 0 || tmp_subscripts[i - 1] != c) &&
-          (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c)) {
-        ret[1].append(1, c);
-      }
-    }
-  }
-
-  // Make sure output subscripts are in the input
-  std::bitset<128> input_subscripts_set = str2set(ret[0]);
-  for (const char& c : ret[1]) {
-    CHECK(input_subscripts_set.test(c))
-      << "Output character " << c
-      << " did not appear in the input";
-  }
-
-  // Make sure number operands is equivalent to the number of terms
-  CHECK_EQ(std::count(ret[0].begin(), ret[0].end(), ',') + 1, operands.size())
-    << "Number of einsum subscripts must be equal to the "
-    << "number of operands.";
-
-  return ret;
-}
-
-inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts, 
-                                        const std::vector<Array<PrimExpr>>& operands) {
-  // Parsing
-  std::vector<std::string> parsed_subscripts = _parse_einsum_input(subscripts, operands);
-
-  // Build a few useful list and sets
-  std::vector<std::string> input_list = split(parsed_subscripts[0], ",");
-  size_t isize = input_list.size();
-
-  // Get length of each unique dimension and ensure all dimensions are correct
-  int dimension_dict[128];
-  memset(dimension_dict, -1, sizeof(dimension_dict));
-  for (size_t i = 0; i < isize; ++i) {
-    const std::string& term = input_list[i];
-    const Array<PrimExpr>& sh = operands[i];
-    CHECK_EQ(sh.size(), term.length())
-      << "Einstein sum subscript " << input_list[i]
-      << " does not contain the "
-      << "correct number of indices for operand " << i << ".";
-    size_t len_term = term.length();
-    for (size_t j = 0; j < len_term; ++j) {
-      int64_t dim = GetConstInt(sh[j]);
-      const char& c = term[j];
-
-      if (dimension_dict[static_cast<int>(c)] != -1) {
-        // For broadcasting cases we always want the largest dim size
-        if (dimension_dict[static_cast<int>(c)] == 1) {
-          dimension_dict[static_cast<int>(c)] = dim;
-        }
-        CHECK(dim == 1 || dim == dimension_dict[static_cast<int>(c)])
-          << "Size of label '" << c
-          << "' for operand  " << i
-          << " (" << dimension_dict[static_cast<int>(c)]
-          << ") does not match previous terms ("
-          << dim << ").";
-      } else {
-        dimension_dict[static_cast<int>(c)] = dim;
-      }
-    }
-  }
-
-  // Get oshape
-  const std::string& output_str = parsed_subscripts[1];
-  size_t odim = output_str.size();
-  Array<PrimExpr> oshape(odim, -1);
-  for (size_t i = 0; i < odim; ++i) {
-    oshape.Set(i, dimension_dict[static_cast<int>(output_str[i])]);
-  }
-  // Neglecting oshape assign check temporally
-  return oshape;
-}
-
-/*!
- * \brief Evaluates the Einstein summation convention on the operands.
- *
- * \param subscripts Specifies the subscripts for summation as comma separated list of subscript labels.
- * \param inputs Arrays for the operation
- * \param name The name of the operation
- * \param tag The tag to mark the operation
- *
- * \return The calculation based on the Einstein summation convention.
- */
-inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inputs,
-                     std::string name = "T_einsum", std::string tag = kMatMul) {
-  // able to compute op: trace, diag, sum, transpose, matmul, dot, inner, outer, multiply, tensordot
-  bool back = false;
-  const char* subscripts = subscripts_str.data();
-  const int nop = inputs.size();
-
-  /* Step 1: Parse the subscripts string into label_counts and op_labels */
-  int iop, idim, min_label = 127, max_label = 0;
-  char label_counts[128], op_labels[16][16];
-  memset(label_counts, 0, sizeof(label_counts));
-  for (iop = 0; iop < nop; ++iop) {
-    int length = static_cast<int>(strcspn(subscripts, ",-"));
-
-    CHECK(!(iop == nop - 1 && subscripts[length] == ','))
-      << "more operands provided to einstein sum function "
-      << "than specified in the subscripts string";
-    CHECK(!(iop < nop-1 && subscripts[length] != ','))
-      << "fewer operands provided to einstein sum function "
-      << "than specified in the subscripts string";
-    parse_operand_subscripts(subscripts, length,
-                                      inputs[iop + back].ndim(),
-                                      iop, op_labels[iop], label_counts,
-                                      &min_label, &max_label);
-
-    /* Move subscripts to the start of the labels for the next op */
-    subscripts += length;
-    if (iop < nop - 1) {
-      subscripts++;
-    }
-  }
-  /*
-   * Find the number of broadcast dimensions, which is the maximum
-   * number of labels == 0 in an op_labels array.
-   */
-  int ndim_broadcast = 0;
-  for (iop = 0; iop < nop; ++iop) {
-    int count_zeros = 0;
-    int ndim;
-    char *labels = op_labels[iop];
-
-    ndim = inputs[iop + back].ndim();
-    for (idim = 0; idim < ndim; ++idim) {
-      if (labels[idim] == 0) {
-        ++count_zeros;
-      }
-    }
-
-    if (count_zeros > ndim_broadcast) {
-      ndim_broadcast = count_zeros;
-    }
-  }
-
-  /*
-   * If there is no output signature, fill output_labels and ndim_output
-   * using each label that appeared once, in alphabetical order.
-   */
-  int label, ndim_output;
-  char output_labels[16];
-  if (subscripts[0] == '\0') {
-    /* If no output was specified, always broadcast left, as usual. */
-    for (ndim_output = 0; ndim_output < ndim_broadcast; ++ndim_output) {
-      output_labels[ndim_output] = 0;
-    }
-    for (label = min_label; label <= max_label; ++label) {
-      if (label_counts[label] == 1) {
-        CHECK(ndim_output < 16)
-          << "einstein sum subscript string has too many "
-          << "distinct labels";
-        output_labels[ndim_output++] = label;
-      }
-    }
-  } else {
-    CHECK(subscripts[0] == '-' && subscripts[1] == '>')
-      << "einstein sum subscript string does not "
-      << "contain proper '->' output specified";
-    subscripts += 2;
-
-    /* Parse the output subscript string. */
-    ndim_output = parse_output_subscripts(subscripts, strlen(subscripts),
-                                          ndim_broadcast, label_counts,
-                                          output_labels);
-    CHECK_GE(ndim_output, 0);
-  }
-
-  /*
-   * Step 2:
-   * Process all the input ops, combining dimensions into their
-   * diagonal where specified.
-   */
-  std::vector<Array<PrimExpr>> opshape(nop), opstride_true(nop);
-  for (iop = 0; iop < nop; ++iop) {
-    char *labels = op_labels[iop];
-    int combine, ndim;
-
-    ndim = inputs[iop + back].ndim();
-
-    /*
-     * Check whether any dimensions need to be combined
-     *
-     * The char type may be either signed or unsigned, we
-     * need it to be signed here.
-     */
-    combine = 0;
-    for (idim = 0; idim < ndim; ++idim) {
-      if ((signed char)labels[idim] < 0) {
-        combine++;
-      }
-    }
-    /* If any dimensions are combined, create a view which combines them */
-    if (combine) {
-      Array<PrimExpr> tshape(static_cast<size_t>(ndim - combine), -1);
-      Array<PrimExpr> tstride(static_cast<size_t>(ndim - combine), -1);
-      GetCombinedDimsView(inputs[iop + back], iop, labels,
-                             tshape, tstride);
-      opshape[iop] = tshape;
-      opstride_true[iop] = tstride;
-    } else {
-      /* No combining needed */
-      opshape[iop] = inputs[iop + back]->shape;
-      opstride_true[iop] = get_stride(opshape[iop]);
-    }
-  }
-  /*
-   * Step 3:
-   * Set up the labels for the iterator (output + combined labels).
-   * Can just share the output_labels memory, because iter_labels
-   * is output_labels with some more labels appended.
-   */
-  char *iter_labels = output_labels;
-  int ndim_iter = ndim_output;
-  for (label = min_label; label <= max_label; ++label) {
-    if (label_counts[label] > 0 &&
-        memchr(output_labels, label, ndim_output) == nullptr) {
-      CHECK(ndim_iter < 16)
-        << "too many subscripts in einsum";
-      iter_labels[ndim_iter++] = label;
-    }
-  }
-  /* Step 4: Set up the op_axes for the iterator */
-  Array<PrimExpr> itershape(static_cast<size_t>(ndim_iter), -1);
-  std::vector<Array<PrimExpr>> iterstride(nop + 1, 
-                                Array<PrimExpr>(static_cast<size_t>(ndim_iter), 0));
-
-  // output_shape
-  std::vector<Array<PrimExpr>> operands;
-  for (size_t i = 0; i < inputs.size(); i++) {
-    operands.push_back(inputs[i]->shape);
-  }
-  Array<PrimExpr> oshape = NumpyEinsumShape(subscripts_str, operands);
-  Array<PrimExpr> ostride_true = get_stride(oshape);
-  Array<PrimExpr> reduceshape;
-  std::vector<Array<PrimExpr>> remainshape(nop);
-  int op_axes_arrays[16][16];
-  int *op_axes[16];
-  for (iop = 0; iop < nop; ++iop) {
-    op_axes[iop] = op_axes_arrays[iop];
-    CHECK_GE(prepare_op_axes(opshape[iop].size(), iop, op_labels[iop],
-             op_axes[iop], ndim_iter, iter_labels), 0);
-    for (idim = 0; idim < ndim_iter; idim++) {
-      if (op_axes[iop][idim] != -1) {
-        iterstride[iop].Set(idim, opstride_true[iop][op_axes[iop][idim]]);
-        if (GetConstInt(itershape[idim]) != -1) {
-          if (GetConstInt(itershape[idim]) == 1) {
-            itershape.Set(idim, opshape[iop][op_axes[iop][idim]]);
-          }
-        } else {
-          itershape.Set(idim, opshape[iop][op_axes[iop][idim]]);
-        }
-      }
-    }
-  }
-  for (idim = 0; idim < ndim_output; ++idim) {
-    iterstride[nop].Set(idim, ostride_true[idim]);
-  }
-  reduceshape = Array<PrimExpr>(static_cast<size_t>(ndim_iter - ndim_output), 0);
-  for (idim = ndim_output; idim < ndim_iter; ++idim) {
-    reduceshape.Set(idim - ndim_output, itershape[idim]);
-  }
-  for (iop = 0; iop < nop; iop++) {
-    Array<Integer> rsh;
-    for (idim = 0; idim < ndim_iter; idim++) {
-      if (op_axes_arrays[iop][idim] == -1) {
-        rsh.push_back(GetConstInt(itershape[idim]));
-      } else {
-        if (GetConstInt(itershape[idim] != opshape[iop][op_axes_arrays[iop][idim]])) {
-          rsh.push_back(GetConstInt(itershape[idim]));
-        }
-      }
-    }
-    remainshape[iop] = Array<PrimExpr>(rsh.begin(), rsh.end());
-  }
-  // exclude the 0-dim case
-  if (ndim_iter == 0) {
-    ndim_iter = 1;
-  }
-  itershape = pad(itershape, ndim_iter);
-  for (iop = 0; iop <= nop; ++iop) {
-    iterstride[iop] = pad(iterstride[iop], ndim_iter);
-  }
-  // oshape = pad(oshape, ndim_iter);
-  reduceshape = pad(reduceshape, ndim_iter);
-  for (iop = 0; iop < nop; ++iop) {
-    opshape[iop] = pad(opshape[iop], ndim_iter);
-    remainshape[iop] = pad(remainshape[iop], ndim_iter);
-  }
-  // ostride and rstride
-  Array<Array<PrimExpr>> ostride;
-  Array<Array<PrimExpr>> rstride;
-
-  for (iop = 0; iop < nop; ++iop) {
-    Array<PrimExpr> otmp(static_cast<size_t>(ndim_iter), 0);
-    Array<PrimExpr> rtmp(static_cast<size_t>(ndim_iter), 0);
-    for (idim = 0; idim < ndim_iter; ++idim) {
-      otmp.Set(idim, idim < ndim_output ? iterstride[iop][idim] : 1);
-      rtmp.Set(idim, idim < ndim_iter - ndim_output ? iterstride[iop][idim+ndim_output] : 1);
-    }
-    ostride.push_back(otmp);
-    rstride.push_back(rtmp);
-  }
-
-
-  // func: input indices => return cooresponding value
-  auto func = [inputs, oshape, ostride, reduceshape, ndim_iter,
-               rstride, nop](const Array<Var>& input_indices) -> PrimExpr {
-    for (int rdim = 0; rdim < ndim_iter; ++rdim) {
-      if (GetConstInt(reduceshape[rdim]) == 0) {
-        return 0;  //
-      }
-    }
-    Array<PrimExpr> ridx = UnravelIndex(0, reduceshape);
-    // AType needs to fit the output data
-    PrimExpr sum = 0;
-    bool rec_flag = false;
-    do {
-      PrimExpr tmp = 1;
-      for (int iop = 0; iop < nop; ++iop) {
-        if (iop != -1) {
-            PrimExpr k = 0;
-            for (size_t i = 0; i < input_indices.size(); ++i) {
-              k += input_indices[i] * ostride[iop][i];
-            }
-            for (size_t i = 0; i < ridx.size(); ++i) {
-              k += ridx[i] * rstride[iop][i];
-            }
-           Array<PrimExpr> temp_indices = UnravelIndex(k, inputs[iop]->shape);
-           tmp = tmp * inputs[iop](temp_indices);
-        }
-      }
-      sum += tmp;
-      
-      ridx.Set(ridx.size() - 1, ridx[ridx.size()-1] + 1);
-      for (int i = static_cast<int>(ridx.size() - 1); (i > 0) &&
-                                  GetConstInt(ridx[i] >= reduceshape[i]); --i) {
-        ridx.Set(i, ridx[i] - reduceshape[i]);
-        ridx.Set(i-1, ridx[i-1] + 1);
-      }
-      rec_flag = GetConstInt(ridx[0] < reduceshape[0]);
-    } while (rec_flag);
-    return sum;
-  };
-
-  return compute(oshape, func, name, tag);
-}
-
 inline Tensor arange(const PrimExpr& start, const PrimExpr& stop, const PrimExpr& step,
                      DataType dtype, std::string name = "T_arange", std::string tag = kInjective) {
   PrimExpr num_elem = tvm::cast(
diff --git a/src/topi/transform.cc b/src/topi/transform.cc
index af47326eba36..17624ed0a0c6 100644
--- a/src/topi/transform.cc
+++ b/src/topi/transform.cc
@@ -24,6 +24,7 @@
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/registry.h>
 #include <tvm/topi/transform.h>
+#include <tvm/topi/einsum.h>
 #include <tvm/topi/utils.h>
 
 namespace tvm {

From ca88de392d07e01c886be594b7b6fa88865bed22 Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 2 Sep 2020 20:15:55 +0800
Subject: [PATCH 06/14] * Fix single value situation

---
 python/tvm/topi/transform.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 35075189d0d7..1bea6696c13a 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -541,6 +541,8 @@ def tensordot(a, b, axes):
 def einsum(subscripts, a_tuple):
     # Grab non-einsum kwargs; do not optimize by default.
     # Temporally not implement the optimized einsum
+    if(not isinstance(a_tuple, tuple)):
+        a_tuple = (a_tuple,)
     return cpp.einsum(subscripts, a_tuple)
 
 

From 80ed874d76012ca45a6bcf4a24d71599c05d4da1 Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 2 Sep 2020 21:23:21 +0800
Subject: [PATCH 07/14] * Fix CamelASE

---
 include/tvm/topi/einsum.h | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/include/tvm/topi/einsum.h b/include/tvm/topi/einsum.h
index 7d797b96fbbd..07077532da07 100644
--- a/include/tvm/topi/einsum.h
+++ b/include/tvm/topi/einsum.h
@@ -44,7 +44,7 @@ namespace topi {
 using namespace tvm::te;
 using namespace topi::detail;
 
-inline Array<PrimExpr> get_stride(const Array<PrimExpr> shape) {
+inline Array<PrimExpr> GetStride(const Array<PrimExpr> shape) {
   size_t ndim = shape.size();
   int prod = 1;
   Array<PrimExpr> stride = Array<PrimExpr>(ndim, -1);
@@ -65,7 +65,7 @@ inline Array<PrimExpr> pad(const Array<PrimExpr> shape, int odim) {
   return ret;
 }
 
-inline int parse_operand_subscripts(const char *subscripts, int length,
+inline int ParseOperandSubscripts(const char *subscripts, int length,
                                     int ndim, int iop, char *op_labels,
                                     char *label_counts, int *min_label, int *max_label) {
   int i;
@@ -156,7 +156,7 @@ inline int parse_operand_subscripts(const char *subscripts, int length,
   return 0;
 }
 
-inline int parse_output_subscripts(const char *subscripts, int length,
+inline int ParseOutputSubscripts(const char *subscripts, int length,
                                    int ndim_broadcast,
                                    const char *label_counts, char *out_labels) {
   int i, bdim;
@@ -233,7 +233,7 @@ inline void GetCombinedDimsView(const Tensor& op, int iop,
   int newdim;
 
   Array<PrimExpr> shape = op->shape;
-  Array<PrimExpr> stride = get_stride(shape);
+  Array<PrimExpr> stride = GetStride(shape);
   ndim = op.ndim();
   newdim = newshape.size();
 
@@ -285,7 +285,7 @@ inline void GetCombinedDimsView(const Tensor& op, int iop,
   }
 }
 
-inline static int prepare_op_axes(int ndim, int iop, char *labels,
+inline static int PrepareOpAxes(int ndim, int iop, char *labels,
                                   int *axes, int ndim_iter, char *iter_labels) {
   int i, label, ibroadcast;
 
@@ -358,7 +358,7 @@ inline std::vector<std::string> split(const std::string& str,
   return ret;
 }
 
-inline std::vector<std::string> _parse_einsum_input(
+inline std::vector<std::string> ParseEinsumInput(
   std::string subscripts,
   const std::vector<Array<PrimExpr>>& operands) {
   const std::string einsum_symbols =
@@ -543,7 +543,7 @@ inline std::vector<std::string> _parse_einsum_input(
 inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
                                         const std::vector<Array<PrimExpr>>& operands) {
   // Parsing
-  std::vector<std::string> parsed_subscripts = _parse_einsum_input(subscripts, operands);
+  std::vector<std::string> parsed_subscripts = ParseEinsumInput(subscripts, operands);
 
   // Build a few useful list and sets
   std::vector<std::string> input_list = split(parsed_subscripts[0], ",");
@@ -622,7 +622,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     CHECK(!(iop < nop-1 && subscripts[length] != ','))
       << "fewer operands provided to einstein sum function "
       << "than specified in the subscripts string";
-    parse_operand_subscripts(subscripts, length,
+    ParseOperandSubscripts(subscripts, length,
                                       inputs[iop + back].ndim(),
                                       iop, op_labels[iop], label_counts,
                                       &min_label, &max_label);
@@ -681,7 +681,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     subscripts += 2;
 
     /* Parse the output subscript string. */
-    ndim_output = parse_output_subscripts(subscripts, strlen(subscripts),
+    ndim_output = ParseOutputSubscripts(subscripts, strlen(subscripts),
                                           ndim_broadcast, label_counts,
                                           output_labels);
     CHECK_GE(ndim_output, 0);
@@ -722,7 +722,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     } else {
       /* No combining needed */
       opshape[iop] = inputs[iop + back]->shape;
-      opstride_true[iop] = get_stride(opshape[iop]);
+      opstride_true[iop] = GetStride(opshape[iop]);
     }
   }
   /*
@@ -752,14 +752,14 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     operands.push_back(inputs[i]->shape);
   }
   Array<PrimExpr> oshape = NumpyEinsumShape(subscripts_str, operands);
-  Array<PrimExpr> ostride_true = get_stride(oshape);
+  Array<PrimExpr> ostride_true = GetStride(oshape);
   Array<PrimExpr> reduceshape;
   std::vector<Array<PrimExpr>> remainshape(nop);
   int op_axes_arrays[16][16];
   int *op_axes[16];
   for (iop = 0; iop < nop; ++iop) {
     op_axes[iop] = op_axes_arrays[iop];
-    CHECK_GE(prepare_op_axes(opshape[iop].size(), iop, op_labels[iop],
+    CHECK_GE(PrepareOpAxes(opshape[iop].size(), iop, op_labels[iop],
              op_axes[iop], ndim_iter, iter_labels), 0);
     for (idim = 0; idim < ndim_iter; idim++) {
       if (op_axes[iop][idim] != -1) {

From bfa7109d1486468b391264f87fb952e81fd7c5ec Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 2 Sep 2020 22:19:10 +0800
Subject: [PATCH 08/14] * Print stride

---
 include/tvm/topi/einsum.h | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/include/tvm/topi/einsum.h b/include/tvm/topi/einsum.h
index 07077532da07..3397a8189e99 100644
--- a/include/tvm/topi/einsum.h
+++ b/include/tvm/topi/einsum.h
@@ -823,6 +823,24 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     rstride.push_back(rtmp);
   }
 
+  // Testing Code for single input
+  for (iop = 0; iop < nop; ++iop) {
+    for(size_t i = 0; i<ostride.size(); i++) {
+      std::printf("ostride[%d][%ld] = %ld\n", iop, i, GetConstInt(ostride[iop][i]));
+    }
+  }
+  for (iop = 0; iop < nop; ++iop) {
+    for(size_t i = 0; i<rstride.size(); i++) {
+      std::printf("rstride[%d][%ld] = %ld\n", iop, i, GetConstInt(rstride[iop][i]));
+    }
+  }
+
+  for (size_t i = 0; i<reduceshape.size(); i++) {
+      std::printf("reduceshape[%ld] = [%ld]\n", i, GetConstInt(reduceshape[i]));
+  }
+
+  
+  // Testing Code End
 
   // func: input indices => return cooresponding value
   auto func = [inputs, oshape, ostride, reduceshape, ndim_iter,
@@ -841,6 +859,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
       for (int iop = 0; iop < nop; ++iop) {
         if (iop != -1) {
             PrimExpr k = 0;
+            // k = dot(input_indices, ostride) + dot(ridx, rstride)
             for (size_t i = 0; i < input_indices.size(); ++i) {
               k += input_indices[i] * ostride[iop][i];
             }

From 33a11daeab0351078633e718e6aa45c998b3bea4 Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 2 Sep 2020 23:45:43 +0800
Subject: [PATCH 09/14] * Fix single input bug

---
 include/tvm/topi/einsum.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/include/tvm/topi/einsum.h b/include/tvm/topi/einsum.h
index 3397a8189e99..26f5ec9827cc 100644
--- a/include/tvm/topi/einsum.h
+++ b/include/tvm/topi/einsum.h
@@ -226,8 +226,8 @@ inline int ParseOutputSubscripts(const char *subscripts, int length,
 
 inline void GetCombinedDimsView(const Tensor& op, int iop,
                                    char *labels,
-                                   Array<PrimExpr> newshape,
-                                   Array<PrimExpr> newstride) {
+                                   Array<PrimExpr> &newshape,
+                                   Array<PrimExpr> &newstride) {
   int idim, ndim, icombine, combineoffset;
   int icombinemap[16];
   int newdim;

From da6edfcb09eb4985f68a1496500f48282e2c80df Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Thu, 3 Sep 2020 00:13:24 +0800
Subject: [PATCH 10/14] * fix lint

---
 include/tvm/topi/einsum.h | 44 ++++++++++++---------------------------
 1 file changed, 13 insertions(+), 31 deletions(-)

diff --git a/include/tvm/topi/einsum.h b/include/tvm/topi/einsum.h
index 26f5ec9827cc..d74d30269856 100644
--- a/include/tvm/topi/einsum.h
+++ b/include/tvm/topi/einsum.h
@@ -226,8 +226,8 @@ inline int ParseOutputSubscripts(const char *subscripts, int length,
 
 inline void GetCombinedDimsView(const Tensor& op, int iop,
                                    char *labels,
-                                   Array<PrimExpr> &newshape,
-                                   Array<PrimExpr> &newstride) {
+                                   Array<PrimExpr>* newshape,
+                                   Array<PrimExpr>* newstride) {
   int idim, ndim, icombine, combineoffset;
   int icombinemap[16];
   int newdim;
@@ -235,12 +235,12 @@ inline void GetCombinedDimsView(const Tensor& op, int iop,
   Array<PrimExpr> shape = op->shape;
   Array<PrimExpr> stride = GetStride(shape);
   ndim = op.ndim();
-  newdim = newshape.size();
+  newdim = newshape->size();
 
   /* Initialize the dimensions and strides to zero */
   for (idim = 0; idim < newdim; ++idim) {
-    newshape.Set(idim, 0);
-    newstride.Set(idim, 0);
+    newshape->Set(idim, 0);
+    newstride->Set(idim, 0);
   }
 
   /* Copy the dimensions and strides, except when collapsing */
@@ -264,18 +264,19 @@ inline void GetCombinedDimsView(const Tensor& op, int iop,
     }
     /* If the label is 0, it's an unlabeled broadcast dimension */
     if (label == 0) {
-      newshape.Set(icombine, shape[idim]);
-      newstride.Set(icombine, stride[idim]);
+      newshape->Set(icombine, shape[idim]);
+      newstride->Set(icombine, stride[idim]);
     } else {
       /* Update the combined axis dimensions and strides */
       int i = icombinemap[idim + combineoffset];
-      CHECK(!((combineoffset < 0) && GetConstInt(newshape[i] != 0 && newshape[i] != shape[idim])))
+      CHECK(!((combineoffset < 0) && GetConstInt((*newshape)[i] != 0 &&
+                                         (*newshape)[i] != shape[idim])))
         << "dimensions in operand " << iop
         << " for collapsing index '" << label
-        << "' don't match (" << GetConstInt(newshape[i])
+        << "' don't match (" << GetConstInt((*newshape)[i])
         << " != " << shape[idim] << ")";
-      newshape.Set(i, shape[idim]);
-      newstride.Set(i, newstride[i] + stride[idim]);
+      newshape->Set(i, shape[idim]);
+      newstride->Set(i, (*newstride)[i] + stride[idim]);
     }
 
     /* If the label didn't say to combine axes, increment dest i */
@@ -716,7 +717,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
       Array<PrimExpr> tshape(static_cast<size_t>(ndim - combine), -1);
       Array<PrimExpr> tstride(static_cast<size_t>(ndim - combine), -1);
       GetCombinedDimsView(inputs[iop + back], iop, labels,
-                             tshape, tstride);
+                             &tshape, &tstride);
       opshape[iop] = tshape;
       opstride_true[iop] = tstride;
     } else {
@@ -823,25 +824,6 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     rstride.push_back(rtmp);
   }
 
-  // Testing Code for single input
-  for (iop = 0; iop < nop; ++iop) {
-    for(size_t i = 0; i<ostride.size(); i++) {
-      std::printf("ostride[%d][%ld] = %ld\n", iop, i, GetConstInt(ostride[iop][i]));
-    }
-  }
-  for (iop = 0; iop < nop; ++iop) {
-    for(size_t i = 0; i<rstride.size(); i++) {
-      std::printf("rstride[%d][%ld] = %ld\n", iop, i, GetConstInt(rstride[iop][i]));
-    }
-  }
-
-  for (size_t i = 0; i<reduceshape.size(); i++) {
-      std::printf("reduceshape[%ld] = [%ld]\n", i, GetConstInt(reduceshape[i]));
-  }
-
-  
-  // Testing Code End
-
   // func: input indices => return cooresponding value
   auto func = [inputs, oshape, ostride, reduceshape, ndim_iter,
                rstride, nop](const Array<Var>& input_indices) -> PrimExpr {

From dfe53709d18e058b5f36cfd553f96ded57099154 Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Wed, 9 Sep 2020 01:46:10 +0800
Subject: [PATCH 11/14] * Fix lint and add comments

---
 include/tvm/topi/einsum.h                 | 434 ++++++++++++----------
 include/tvm/topi/transform.h              |   1 -
 python/tvm/topi/transform.py              |   2 +-
 src/topi/transform.cc                     |   2 +-
 tests/python/unittest/test_te_autodiff.py |   4 +
 5 files changed, 252 insertions(+), 191 deletions(-)

diff --git a/include/tvm/topi/einsum.h b/include/tvm/topi/einsum.h
index d74d30269856..29df13460a7f 100644
--- a/include/tvm/topi/einsum.h
+++ b/include/tvm/topi/einsum.h
@@ -32,11 +32,11 @@
 #include <tvm/topi/tags.h>
 
 #include <algorithm>
+#include <bitset>
 #include <iterator>
 #include <string>
 #include <unordered_set>
 #include <vector>
-#include <bitset>
 
 namespace tvm {
 namespace topi {
@@ -44,6 +44,13 @@ namespace topi {
 using namespace tvm::te;
 using namespace topi::detail;
 
+/*!
+ * \brief Compute the stride of the given shape.
+ *
+ * \param shape for the operation.
+ *
+ * \return the stride of the shape.
+ */
 inline Array<PrimExpr> GetStride(const Array<PrimExpr> shape) {
   size_t ndim = shape.size();
   int prod = 1;
@@ -55,7 +62,15 @@ inline Array<PrimExpr> GetStride(const Array<PrimExpr> shape) {
   return stride;
 }
 
-inline Array<PrimExpr> pad(const Array<PrimExpr> shape, int odim) {
+/*!
+ * \brief Pad the shape with 1.
+ *
+ * \param shape the input shape to be padded
+ * \param odim the padding size of the objective shape.
+ *
+ * \return the padded shape.
+ */
+inline Array<PrimExpr> Pad(const Array<PrimExpr> shape, int odim) {
   int ndim = shape.size();
   CHECK_GE(odim, ndim);
   Array<PrimExpr> ret(static_cast<size_t>(odim), 1);
@@ -65,9 +80,26 @@ inline Array<PrimExpr> pad(const Array<PrimExpr> shape, int odim) {
   return ret;
 }
 
-inline int ParseOperandSubscripts(const char *subscripts, int length,
-                                    int ndim, int iop, char *op_labels,
-                                    char *label_counts, int *min_label, int *max_label) {
+/*!
+ * \brief Parse the subscripts for one operand into an output of 'ndim' labels.
+ *
+ * \param subscripts the subscripts for to be parsed.
+ * \param length subscripts[0: length] represents the current operand.
+ * \param ndim the ndim of current operand.
+ * \param iop the index of the operand.
+ * \param op_labels the parsing result.
+ *        For Example:
+ *           subscripts="abbcbc",  ndim=6 -> op_labels=[97, 98, -1, 99, -3, -2].
+ *           subscripts="ab...bc", ndim=6 -> op_labels=[97, 98, 0, 0, -3, 99].
+ * \param label_counts Count the number the label appears.
+ * \param min_label Save the minimal label according to ASCII.
+ * \param max_label Save the maximal label according to ASCII.
+ *
+ * \return 0.
+ */
+inline int ParseOperandSubscripts(const char* subscripts, int length, int ndim, int iop,
+                                  char* op_labels, char* label_counts, int* min_label,
+                                  int* max_label) {
   int i;
   int idim = 0;
   int ellipsis = -1;
@@ -79,10 +111,8 @@ inline int ParseOperandSubscripts(const char *subscripts, int length,
     /* A proper label for an axis. */
     if (label > 0 && isalpha(label)) {
       /* Check we don't exceed the operator dimensions. */
-      CHECK(idim < ndim)
-        << "einstein sum subscripts string contains "
-        << "too many subscripts for operand "
-        << iop;
+      CHECK(idim < ndim) << "einstein sum subscripts string contains "
+                         << "too many subscripts for operand " << iop;
 
       op_labels[idim++] = label;
       if (label < *min_label) {
@@ -95,29 +125,26 @@ inline int ParseOperandSubscripts(const char *subscripts, int length,
     } else if (label == '.') {
       /* The beginning of the ellipsis. */
       /* Check it's a proper ellipsis. */
-      CHECK(!(ellipsis != -1 || i + 2 >= length
-              || subscripts[++i] != '.' || subscripts[++i] != '.'))
-        << "einstein sum subscripts string contains a "
-        << "'.' that is not part of an ellipsis ('...') "
-        << "in operand "
-        << iop;
+      CHECK(
+          !(ellipsis != -1 || i + 2 >= length || subscripts[++i] != '.' || subscripts[++i] != '.'))
+          << "einstein sum subscripts string contains a "
+          << "'.' that is not part of an ellipsis ('...') "
+          << "in operand " << iop;
 
       ellipsis = idim;
     } else {
-        CHECK(label == ' ')
-          << "invalid subscript '" << static_cast<char>(label)
-          << "' in einstein sum "
-          << "subscripts string, subscripts must "
-          << "be letters";
+      CHECK(label == ' ') << "invalid subscript '" << static_cast<char>(label)
+                          << "' in einstein sum "
+                          << "subscripts string, subscripts must "
+                          << "be letters";
     }
   }
 
   /* No ellipsis found, labels must match dimensions exactly. */
   if (ellipsis == -1) {
-    CHECK(idim == ndim)
-      << "operand has more dimensions than subscripts "
-      << "given in einstein sum, but no '...' ellipsis "
-      << "provided to broadcast the extra dimensions.";
+    CHECK(idim == ndim) << "operand has more dimensions than subscripts "
+                        << "given in einstein sum, but no '...' ellipsis "
+                        << "provided to broadcast the extra dimensions.";
   } else if (idim < ndim) {
     /* Ellipsis found, may have to add broadcast dimensions. */
     /* Move labels after ellipsis to the end. */
@@ -143,7 +170,7 @@ inline int ParseOperandSubscripts(const char *subscripts, int length,
     /* If it is a proper label, find any duplicates of it. */
     if (label > 0) {
       /* Search for the next matching label. */
-      char *next = reinterpret_cast<char*>(memchr(op_labels + idim + 1, label, ndim - idim - 1));
+      char* next = reinterpret_cast<char*>(memchr(op_labels + idim + 1, label, ndim - idim - 1));
 
       while (next != nullptr) {
         /* The offset from next to op_labels[idim] (negative). */
@@ -156,9 +183,21 @@ inline int ParseOperandSubscripts(const char *subscripts, int length,
   return 0;
 }
 
-inline int ParseOutputSubscripts(const char *subscripts, int length,
-                                   int ndim_broadcast,
-                                   const char *label_counts, char *out_labels) {
+/*!
+ * \brief Parse the subscripts for the output into an output that includes 'ndim_broadcast'
+ *        unlabeled dimensions.
+ *
+ * \param subscripts the subscripts for to be parsed.
+ * \param length subscripts[0: length] represents the output operand.
+ * \param ndim_broadcast the broadcast dimension number.
+ * \param label_counts Count the number the label appears.
+ * \param out_labels similar to the op_labels in ParseOperandSubscripts, for each
+ *        dimension, the ASCII code of the corresponding label. zero for the broadcasting dim.
+ *
+ * \return the total number of output dimensions or -1 if there is an error.
+ */
+inline int ParseOutputSubscripts(const char* subscripts, int length, int ndim_broadcast,
+                                 const char* label_counts, char* out_labels) {
   int i, bdim;
   int ndim = 0;
   int ellipsis = 0;
@@ -171,63 +210,68 @@ inline int ParseOutputSubscripts(const char *subscripts, int length,
     if (label > 0 && isalpha(label)) {
       /* Check that it doesn't occur again. */
       CHECK(memchr(subscripts + i + 1, label, length - i - 1) == nullptr)
-        << "einstein sum subscripts string includes "
-        << "output subscript '" << static_cast<char>(label)
-        << "' multiple times";
+          << "einstein sum subscripts string includes "
+          << "output subscript '" << static_cast<char>(label) << "' multiple times";
 
       /* Check that it was used in the inputs. */
       CHECK(label_counts[label] != 0)
-        << "einstein sum subscripts string included "
-        << "output subscript '" << static_cast<char>(label)
-        << "' which never appeared "
-        << "in an input";
+          << "einstein sum subscripts string included "
+          << "output subscript '" << static_cast<char>(label) << "' which never appeared "
+          << "in an input";
 
       /* Check that there is room in out_labels for this label. */
-      CHECK(ndim < 16)
-        << "einstein sum subscripts string contains "
-        << "too many subscripts in the output";
+      CHECK(ndim < 16) << "einstein sum subscripts string contains "
+                       << "too many subscripts in the output";
 
       out_labels[ndim++] = label;
     } else if (label == '.') {
       /* The beginning of the ellipsis. */
       /* Check it is a proper ellipsis. */
-      CHECK(!(ellipsis || i + 2 >= length
-              || subscripts[++i] != '.' || subscripts[++i] != '.'))
-        << "einstein sum subscripts string "
-        << "contains a '.' that is not part of "
-        << "an ellipsis ('...') in the output";
+      CHECK(!(ellipsis || i + 2 >= length || subscripts[++i] != '.' || subscripts[++i] != '.'))
+          << "einstein sum subscripts string "
+          << "contains a '.' that is not part of "
+          << "an ellipsis ('...') in the output";
 
       /* Check there is room in out_labels for broadcast dims. */
-      CHECK(ndim + ndim_broadcast <= 16)
-        << "einstein sum subscripts string contains "
-        << "too many subscripts in the output";
+      CHECK(ndim + ndim_broadcast <= 16) << "einstein sum subscripts string contains "
+                                         << "too many subscripts in the output";
 
       ellipsis = 1;
       for (bdim = 0; bdim < ndim_broadcast; ++bdim) {
         out_labels[ndim++] = 0;
       }
     } else {
-      CHECK(label == ' ')
-        << "invalid subscript '" << static_cast<char>(label)
-        << "' in einstein sum "
-        << "subscripts string, subscripts must "
-        << "be letters";
+      CHECK(label == ' ') << "invalid subscript '" << static_cast<char>(label)
+                          << "' in einstein sum "
+                          << "subscripts string, subscripts must "
+                          << "be letters";
     }
   }
 
   /* If no ellipsis was found there should be no broadcast dimensions. */
-  CHECK(!(!ellipsis && ndim_broadcast > 0))
-    << "output has more dimensions than subscripts "
-    << "given in einstein sum, but no '...' ellipsis "
-    << "provided to broadcast the extra dimensions.";
+  CHECK(!(!ellipsis && ndim_broadcast > 0)) << "output has more dimensions than subscripts "
+                                            << "given in einstein sum, but no '...' ellipsis "
+                                            << "provided to broadcast the extra dimensions.";
 
   return ndim;
 }
 
-inline void GetCombinedDimsView(const Tensor& op, int iop,
-                                   char *labels,
-                                   Array<PrimExpr>* newshape,
-                                   Array<PrimExpr>* newstride) {
+/*!
+ * \brief If any dimensions are combined, create a view that combines them.
+ *        Shows in newshape and newstride.
+ *
+ * \param op the operand tensor.
+ * \param iop the index of the operand.
+ * \param labels the op_labels fot the operand. Like [97, 98, -2] for "aba".
+ * \param newshape The combined shape.
+ * \param newstride The combined stride.
+ *
+ * For example:
+ *  "aba -> ab",              shape = [2,3,2] stride = [6,2,1]
+ *  op_labels = [97, 98, -2], newshape = [2,3], newstride = [7,2]
+ */
+inline void GetCombinedDimsView(const Tensor& op, int iop, char* labels, Array<PrimExpr>* newshape,
+                                Array<PrimExpr>* newstride) {
   int idim, ndim, icombine, combineoffset;
   int icombinemap[16];
   int newdim;
@@ -254,7 +298,7 @@ inline void GetCombinedDimsView(const Tensor& op, int iop,
     /* If this label says to merge axes, get the actual label */
     if (label < 0) {
       combineoffset = label;
-      label = labels[idim+label];
+      label = labels[idim + label];
     } else {
       combineoffset = 0;
       if (icombine != idim) {
@@ -269,12 +313,10 @@ inline void GetCombinedDimsView(const Tensor& op, int iop,
     } else {
       /* Update the combined axis dimensions and strides */
       int i = icombinemap[idim + combineoffset];
-      CHECK(!((combineoffset < 0) && GetConstInt((*newshape)[i] != 0 &&
-                                         (*newshape)[i] != shape[idim])))
-        << "dimensions in operand " << iop
-        << " for collapsing index '" << label
-        << "' don't match (" << GetConstInt((*newshape)[i])
-        << " != " << shape[idim] << ")";
+      CHECK(!((combineoffset < 0) &&
+              GetConstInt((*newshape)[i] != 0 && (*newshape)[i] != shape[idim])))
+          << "dimensions in operand " << iop << " for collapsing index '" << label
+          << "' don't match (" << GetConstInt((*newshape)[i]) << " != " << shape[idim] << ")";
       newshape->Set(i, shape[idim]);
       newstride->Set(i, (*newstride)[i] + stride[idim]);
     }
@@ -286,12 +328,22 @@ inline void GetCombinedDimsView(const Tensor& op, int iop,
   }
 }
 
-inline static int PrepareOpAxes(int ndim, int iop, char *labels,
-                                  int *axes, int ndim_iter, char *iter_labels) {
+/*!
+ * \brief Prepare the operand axes to match each stride or shape pair.
+ *
+ * \param ndim the ndim of the operand tensor.
+ * \param iop the index of the operand.
+ * \param labels the op_labels fot the operand. [97, 98, -1, 99, -3, -2] for "abbcbc".
+ * \param axes The matched axes to be calculated.
+ * \param ndim_iter the dimension of iterating. Subscripts "ab, bc -> ac" ndim_iter = 3.
+ * \param iter_labels output_labels with the iterating label. ['a', 'c', 'b'] for the case above.
+ */
+inline static int PrepareOpAxes(int ndim, int iop, char* labels, int* axes, int ndim_iter,
+                                char* iter_labels) {
   int i, label, ibroadcast;
 
-  ibroadcast = ndim-1;
-  for (i = ndim_iter-1; i >= 0; --i) {
+  ibroadcast = ndim - 1;
+  for (i = ndim_iter - 1; i >= 0; --i) {
     label = iter_labels[i];
     /*
      * If it's an unlabeled broadcast dimension, choose
@@ -314,7 +366,7 @@ inline static int PrepareOpAxes(int ndim, int iop, char *labels,
       }
     } else {
       /* It's a labeled dimension, find the matching one */
-      char *match = reinterpret_cast<char*>(memchr(labels, label, ndim));
+      char* match = reinterpret_cast<char*>(memchr(labels, label, ndim));
       /* If the op doesn't have the label, broadcast it */
       if (match == nullptr) {
         axes[i] = -1;
@@ -327,8 +379,14 @@ inline static int PrepareOpAxes(int ndim, int iop, char *labels,
   return 0;
 }
 
-inline int _count_substring(const std::string& str,
-                            const std::string& sub) {
+/*!
+ * \brief Count SubString.
+ * \param str the object string
+ * \param sub the pattern string
+ *
+ * \return number of substring
+ */
+inline int CountSubstring(const std::string& str, const std::string& sub) {
   int count = 0;
   std::string::size_type pos = 0;
   while ((pos = str.find(sub, pos)) != std::string::npos) {
@@ -338,7 +396,13 @@ inline int _count_substring(const std::string& str,
   return count;
 }
 
-inline std::bitset<128> str2set(const std::string& str) {
+/*!
+ * \brief Transfer string to.
+ * \param str input string.
+ *
+ * \return bitset.
+ */
+inline std::bitset<128> Str2Set(const std::string& str) {
   std::bitset<128> ret;
   for (const char& c : str) {
     ret.set(static_cast<int>(c));
@@ -346,8 +410,14 @@ inline std::bitset<128> str2set(const std::string& str) {
   return ret;
 }
 
-inline std::vector<std::string> split(const std::string& str,
-                                      const std::string& sub) {
+/*!
+ * \brief Split str according to substring.
+ * \param str input string.
+ * \param sub the split pattern string.
+ *
+ * \return vector contains the splited substring.
+ */
+inline std::vector<std::string> Split(const std::string& str, const std::string& sub) {
   std::string::size_type pos = 0;
   std::string::size_type start = 0;
   std::vector<std::string> ret;
@@ -359,18 +429,24 @@ inline std::vector<std::string> split(const std::string& str,
   return ret;
 }
 
-inline std::vector<std::string> ParseEinsumInput(
-  std::string subscripts,
-  const std::vector<Array<PrimExpr>>& operands) {
-  const std::string einsum_symbols =
-    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
+/*!
+ * \brief Parse the input subscripts into a vector of strings.
+ * \param subscripts input subscripts.
+ * \param operands operand tensors.
+ *
+ * \return vector of strings, vector[0] represents the input part, vector[1] represents the ouput.
+ * if no output, the vector[1] is NULL.
+ * "ab, bc -> ac" => ["ab,bc", "ac"]
+ */
+inline std::vector<std::string> ParseEinsumInput(std::string subscripts,
+                                                 const std::vector<Array<PrimExpr>>& operands) {
+  const std::string einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
   std::bitset<128> einsum_symbols_set;
   for (const char& c : einsum_symbols) {
     einsum_symbols_set.set(c);
   }
 
-  CHECK_NE(operands.size(), 0U)
-    << "No input operands";
+  CHECK_NE(operands.size(), 0U) << "No input operands";
 
   auto end_pos = std::remove(subscripts.begin(), subscripts.end(), ' ');
   subscripts.erase(end_pos, subscripts.end());
@@ -380,32 +456,26 @@ inline std::vector<std::string> ParseEinsumInput(
     if (c == '.' || c == ',' || c == '-' || c == '>') {
       continue;
     }
-    CHECK(einsum_symbols_set.test(c))
-      << "Character " << c
-      << " is not a valid symbol.";
+    CHECK(einsum_symbols_set.test(c)) << "Character " << c << " is not a valid symbol.";
   }
 
   // Check for proper "->"
-  if (subscripts.find('-') != std::string::npos ||
-      subscripts.find('>') != std::string::npos) {
+  if (subscripts.find('-') != std::string::npos || subscripts.find('>') != std::string::npos) {
     bool invalid = (std::count(subscripts.begin(), subscripts.end(), '-') > 1 ||
                     std::count(subscripts.begin(), subscripts.end(), '>') > 1);
-    CHECK(!invalid && _count_substring(subscripts, "->") == 1)
-      << "Subscripts can only contain one '->'.";
+    CHECK(!invalid && CountSubstring(subscripts, "->") == 1)
+        << "Subscripts can only contain one '->'.";
   }
 
   // Parse ellipses
   if (subscripts.find('.') != std::string::npos) {
     std::string used = subscripts;
-    used.erase(std::remove_if(used.begin(),
-                              used.end(),
-                              [](const char& c){return c == '.' ||
-                                                       c == ',' ||
-                                                       c == '-' ||
-                                                       c == '>';}),
-               used.end());
-
-    std::bitset<128> used_set = str2set(used);
+    used.erase(
+        std::remove_if(used.begin(), used.end(),
+                       [](const char& c) { return c == '.' || c == ',' || c == '-' || c == '>'; }),
+        used.end());
+
+    std::bitset<128> used_set = Str2Set(used);
     std::string ellipse_inds = "";
     for (const char& c : einsum_symbols) {
       if (!used_set.test(static_cast<int>(c))) {
@@ -418,13 +488,13 @@ inline std::vector<std::string> ParseEinsumInput(
     bool out_sub;
 
     if (subscripts.find("->") != std::string::npos) {
-      std::vector<std::string> tmp = split(subscripts, "->");
+      std::vector<std::string> tmp = Split(subscripts, "->");
       input_tmp = tmp[0];
       output_sub = tmp[1];
-      split_subscripts = split(input_tmp, ",");
+      split_subscripts = Split(input_tmp, ",");
       out_sub = true;
     } else {
-      split_subscripts = split(subscripts, ",");
+      split_subscripts = Split(subscripts, ",");
       out_sub = false;
     }
 
@@ -433,10 +503,8 @@ inline std::vector<std::string> ParseEinsumInput(
     for (size_t i = 0; i < size_split_subscripts; ++i) {
       const std::string& sub = split_subscripts[i];
       if (sub.find('.') != std::string::npos) {
-        CHECK_EQ(std::count(sub.begin(), sub.end(), '.'), 3)
-          << "Invalid Ellipses";
-        CHECK_EQ(_count_substring(sub, "..."), 1)
-          << "Invalid Ellipses";
+        CHECK_EQ(std::count(sub.begin(), sub.end(), '.'), 3) << "Invalid Ellipses";
+        CHECK_EQ(CountSubstring(sub, "..."), 1) << "Invalid Ellipses";
 
         // Take into account numerical values
         int ellipse_count = 0;
@@ -451,8 +519,7 @@ inline std::vector<std::string> ParseEinsumInput(
           longest = ellipse_count;
         }
 
-        CHECK_GE(ellipse_count, 0)
-          << "Ellipses lengths do not match.";
+        CHECK_GE(ellipse_count, 0) << "Ellipses lengths do not match.";
         if (ellipse_count == 0) {
           split_subscripts[i].erase(sub.find("..."), 3);
         } else {
@@ -477,7 +544,7 @@ inline std::vector<std::string> ParseEinsumInput(
       subscripts += "->" + output_sub;
     } else {
       // Special care for outputless ellipses
-      std::bitset<128> out_ellipse_set = str2set(out_ellipse);
+      std::bitset<128> out_ellipse_set = Str2Set(out_ellipse);
       std::string tmp_subscripts = subscripts, output_subscript = "";
       size_t len_tmp_subscripts = tmp_subscripts.length();
       std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
@@ -486,9 +553,7 @@ inline std::vector<std::string> ParseEinsumInput(
         if (c == ',') {
           continue;
         }
-        CHECK(einsum_symbols_set.test(c))
-          << "Character " << c
-          << " is not a valid symbol.";
+        CHECK(einsum_symbols_set.test(c)) << "Character " << c << " is not a valid symbol.";
         if ((i == 0 || tmp_subscripts[i - 1] != c) &&
             (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c) &&
             !out_ellipse_set.test(c)) {
@@ -502,7 +567,7 @@ inline std::vector<std::string> ParseEinsumInput(
   // Build output string if does not exist
   std::vector<std::string> ret(2);
   if (subscripts.find("->") != std::string::npos) {
-    ret = split(subscripts, "->");
+    ret = Split(subscripts, "->");
   } else {
     ret[0] = subscripts;
     ret[1] = "";
@@ -515,9 +580,7 @@ inline std::vector<std::string> ParseEinsumInput(
       if (c == ',') {
         continue;
       }
-      CHECK(einsum_symbols_set.test(c))
-        << "Character " << c
-        << " is not a valid symbol.";
+      CHECK(einsum_symbols_set.test(c)) << "Character " << c << " is not a valid symbol.";
       if ((i == 0 || tmp_subscripts[i - 1] != c) &&
           (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c)) {
         ret[1].append(1, c);
@@ -526,28 +589,34 @@ inline std::vector<std::string> ParseEinsumInput(
   }
 
   // Make sure output subscripts are in the input
-  std::bitset<128> input_subscripts_set = str2set(ret[0]);
+  std::bitset<128> input_subscripts_set = Str2Set(ret[0]);
   for (const char& c : ret[1]) {
     CHECK(input_subscripts_set.test(c))
-      << "Output character " << c
-      << " did not appear in the input";
+        << "Output character " << c << " did not appear in the input";
   }
 
   // Make sure number operands is equivalent to the number of terms
   CHECK_EQ(std::count(ret[0].begin(), ret[0].end(), ',') + 1, operands.size())
-    << "Number of einsum subscripts must be equal to the "
-    << "number of operands.";
+      << "Number of einsum subscripts must be equal to the "
+      << "number of operands.";
 
   return ret;
 }
 
+/*!
+ * \brief Compute the shape of the output.
+ * \param subscripts input subscripts.
+ * \param operands operand tensors.
+ *
+ * \return the shape of the output.
+ */
 inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
                                         const std::vector<Array<PrimExpr>>& operands) {
   // Parsing
   std::vector<std::string> parsed_subscripts = ParseEinsumInput(subscripts, operands);
 
   // Build a few useful list and sets
-  std::vector<std::string> input_list = split(parsed_subscripts[0], ",");
+  std::vector<std::string> input_list = Split(parsed_subscripts[0], ",");
   size_t isize = input_list.size();
 
   // Get length of each unique dimension and ensure all dimensions are correct
@@ -557,9 +626,8 @@ inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
     const std::string& term = input_list[i];
     const Array<PrimExpr>& sh = operands[i];
     CHECK_EQ(sh.size(), term.length())
-      << "Einstein sum subscript " << input_list[i]
-      << " does not contain the "
-      << "correct number of indices for operand " << i << ".";
+        << "Einstein sum subscript " << input_list[i] << " does not contain the "
+        << "correct number of indices for operand " << i << ".";
     size_t len_term = term.length();
     for (size_t j = 0; j < len_term; ++j) {
       int64_t dim = GetConstInt(sh[j]);
@@ -571,11 +639,9 @@ inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
           dimension_dict[static_cast<int>(c)] = dim;
         }
         CHECK(dim == 1 || dim == dimension_dict[static_cast<int>(c)])
-          << "Size of label '" << c
-          << "' for operand  " << i
-          << " (" << dimension_dict[static_cast<int>(c)]
-          << ") does not match previous terms ("
-          << dim << ").";
+            << "Size of label '" << c << "' for operand  " << i << " ("
+            << dimension_dict[static_cast<int>(c)] << ") does not match previous terms (" << dim
+            << ").";
       } else {
         dimension_dict[static_cast<int>(c)] = dim;
       }
@@ -596,10 +662,9 @@ inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
 /*!
  * \brief Evaluates the Einstein summation convention on the operands.
  *
- * \param subscripts Specifies the subscripts for summation as comma separated list of subscript labels.
- * \param inputs Arrays for the operation
- * \param name The name of the operation
- * \param tag The tag to mark the operation
+ * \param subscripts_str Specifies the subscripts for summation as comma separated list of subscript
+ * labels. \param inputs Arrays for the operation. \param name The name of the operation. \param tag
+ * The tag to mark the operation.
  *
  * \return The calculation based on the Einstein summation convention.
  */
@@ -618,15 +683,13 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     int length = static_cast<int>(strcspn(subscripts, ",-"));
 
     CHECK(!(iop == nop - 1 && subscripts[length] == ','))
-      << "more operands provided to einstein sum function "
-      << "than specified in the subscripts string";
-    CHECK(!(iop < nop-1 && subscripts[length] != ','))
-      << "fewer operands provided to einstein sum function "
-      << "than specified in the subscripts string";
-    ParseOperandSubscripts(subscripts, length,
-                                      inputs[iop + back].ndim(),
-                                      iop, op_labels[iop], label_counts,
-                                      &min_label, &max_label);
+        << "more operands provided to einstein sum function "
+        << "than specified in the subscripts string";
+    CHECK(!(iop < nop - 1 && subscripts[length] != ','))
+        << "fewer operands provided to einstein sum function "
+        << "than specified in the subscripts string";
+    ParseOperandSubscripts(subscripts, length, inputs[iop + back].ndim(), iop, op_labels[iop],
+                           label_counts, &min_label, &max_label);
 
     /* Move subscripts to the start of the labels for the next op */
     subscripts += length;
@@ -642,7 +705,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   for (iop = 0; iop < nop; ++iop) {
     int count_zeros = 0;
     int ndim;
-    char *labels = op_labels[iop];
+    char* labels = op_labels[iop];
 
     ndim = inputs[iop + back].ndim();
     for (idim = 0; idim < ndim; ++idim) {
@@ -669,22 +732,19 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     }
     for (label = min_label; label <= max_label; ++label) {
       if (label_counts[label] == 1) {
-        CHECK(ndim_output < 16)
-          << "einstein sum subscript string has too many "
-          << "distinct labels";
+        CHECK(ndim_output < 16) << "einstein sum subscript string has too many "
+                                << "distinct labels";
         output_labels[ndim_output++] = label;
       }
     }
   } else {
-    CHECK(subscripts[0] == '-' && subscripts[1] == '>')
-      << "einstein sum subscript string does not "
-      << "contain proper '->' output specified";
+    CHECK(subscripts[0] == '-' && subscripts[1] == '>') << "einstein sum subscript string does not "
+                                                        << "contain proper '->' output specified";
     subscripts += 2;
 
     /* Parse the output subscript string. */
-    ndim_output = ParseOutputSubscripts(subscripts, strlen(subscripts),
-                                          ndim_broadcast, label_counts,
-                                          output_labels);
+    ndim_output = ParseOutputSubscripts(subscripts, strlen(subscripts), ndim_broadcast,
+                                        label_counts, output_labels);
     CHECK_GE(ndim_output, 0);
   }
 
@@ -695,7 +755,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
    */
   std::vector<Array<PrimExpr>> opshape(nop), opstride_true(nop);
   for (iop = 0; iop < nop; ++iop) {
-    char *labels = op_labels[iop];
+    char* labels = op_labels[iop];
     int combine, ndim;
 
     ndim = inputs[iop + back].ndim();
@@ -716,8 +776,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     if (combine) {
       Array<PrimExpr> tshape(static_cast<size_t>(ndim - combine), -1);
       Array<PrimExpr> tstride(static_cast<size_t>(ndim - combine), -1);
-      GetCombinedDimsView(inputs[iop + back], iop, labels,
-                             &tshape, &tstride);
+      GetCombinedDimsView(inputs[iop + back], iop, labels, &tshape, &tstride);
       opshape[iop] = tshape;
       opstride_true[iop] = tstride;
     } else {
@@ -732,20 +791,18 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
    * Can just share the output_labels memory, because iter_labels
    * is output_labels with some more labels appended.
    */
-  char *iter_labels = output_labels;
+  char* iter_labels = output_labels;
   int ndim_iter = ndim_output;
   for (label = min_label; label <= max_label; ++label) {
-    if (label_counts[label] > 0 &&
-        memchr(output_labels, label, ndim_output) == nullptr) {
-      CHECK(ndim_iter < 16)
-        << "too many subscripts in einsum";
+    if (label_counts[label] > 0 && memchr(output_labels, label, ndim_output) == nullptr) {
+      CHECK(ndim_iter < 16) << "too many subscripts in einsum";
       iter_labels[ndim_iter++] = label;
     }
   }
   /* Step 4: Set up the op_axes for the iterator */
   Array<PrimExpr> itershape(static_cast<size_t>(ndim_iter), -1);
   std::vector<Array<PrimExpr>> iterstride(nop + 1,
-                                Array<PrimExpr>(static_cast<size_t>(ndim_iter), 0));
+                                          Array<PrimExpr>(static_cast<size_t>(ndim_iter), 0));
 
   // output_shape
   std::vector<Array<PrimExpr>> operands;
@@ -757,11 +814,12 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   Array<PrimExpr> reduceshape;
   std::vector<Array<PrimExpr>> remainshape(nop);
   int op_axes_arrays[16][16];
-  int *op_axes[16];
+  int* op_axes[16];
   for (iop = 0; iop < nop; ++iop) {
     op_axes[iop] = op_axes_arrays[iop];
-    CHECK_GE(PrepareOpAxes(opshape[iop].size(), iop, op_labels[iop],
-             op_axes[iop], ndim_iter, iter_labels), 0);
+    CHECK_GE(PrepareOpAxes(opshape[iop].size(), iop, op_labels[iop], op_axes[iop], ndim_iter,
+                           iter_labels),
+             0);
     for (idim = 0; idim < ndim_iter; idim++) {
       if (op_axes[iop][idim] != -1) {
         iterstride[iop].Set(idim, opstride_true[iop][op_axes[iop][idim]]);
@@ -799,15 +857,15 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   if (ndim_iter == 0) {
     ndim_iter = 1;
   }
-  itershape = pad(itershape, ndim_iter);
+  itershape = Pad(itershape, ndim_iter);
   for (iop = 0; iop <= nop; ++iop) {
-    iterstride[iop] = pad(iterstride[iop], ndim_iter);
+    iterstride[iop] = Pad(iterstride[iop], ndim_iter);
   }
-  // oshape = pad(oshape, ndim_iter);
-  reduceshape = pad(reduceshape, ndim_iter);
+  // oshape = Pad(oshape, ndim_iter);
+  reduceshape = Pad(reduceshape, ndim_iter);
   for (iop = 0; iop < nop; ++iop) {
-    opshape[iop] = pad(opshape[iop], ndim_iter);
-    remainshape[iop] = pad(remainshape[iop], ndim_iter);
+    opshape[iop] = Pad(opshape[iop], ndim_iter);
+    remainshape[iop] = Pad(remainshape[iop], ndim_iter);
   }
   // ostride and rstride
   Array<Array<PrimExpr>> ostride;
@@ -818,15 +876,15 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     Array<PrimExpr> rtmp(static_cast<size_t>(ndim_iter), 0);
     for (idim = 0; idim < ndim_iter; ++idim) {
       otmp.Set(idim, idim < ndim_output ? iterstride[iop][idim] : 1);
-      rtmp.Set(idim, idim < ndim_iter - ndim_output ? iterstride[iop][idim+ndim_output] : 1);
+      rtmp.Set(idim, idim < ndim_iter - ndim_output ? iterstride[iop][idim + ndim_output] : 1);
     }
     ostride.push_back(otmp);
     rstride.push_back(rtmp);
   }
 
   // func: input indices => return cooresponding value
-  auto func = [inputs, oshape, ostride, reduceshape, ndim_iter,
-               rstride, nop](const Array<Var>& input_indices) -> PrimExpr {
+  auto func = [inputs, oshape, ostride, reduceshape, ndim_iter, rstride,
+               nop](const Array<Var>& input_indices) -> PrimExpr {
     for (int rdim = 0; rdim < ndim_iter; ++rdim) {
       if (GetConstInt(reduceshape[rdim]) == 0) {
         return 0;  //
@@ -840,24 +898,24 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
       PrimExpr tmp = 1;
       for (int iop = 0; iop < nop; ++iop) {
         if (iop != -1) {
-            PrimExpr k = 0;
-            // k = dot(input_indices, ostride) + dot(ridx, rstride)
-            for (size_t i = 0; i < input_indices.size(); ++i) {
-              k += input_indices[i] * ostride[iop][i];
-            }
-            for (size_t i = 0; i < ridx.size(); ++i) {
-              k += ridx[i] * rstride[iop][i];
-            }
-           Array<PrimExpr> temp_indices = UnravelIndex(k, inputs[iop]->shape);
-           tmp = tmp * inputs[iop](temp_indices);
+          PrimExpr k = 0;
+          // k = dot(input_indices, ostride) + dot(ridx, rstride)
+          for (size_t i = 0; i < input_indices.size(); ++i) {
+            k += input_indices[i] * ostride[iop][i];
+          }
+          for (size_t i = 0; i < ridx.size(); ++i) {
+            k += ridx[i] * rstride[iop][i];
+          }
+          Array<PrimExpr> temp_indices = UnravelIndex(k, inputs[iop]->shape);
+          tmp = tmp * inputs[iop](temp_indices);
         }
       }
       sum += tmp;
-      ridx.Set(ridx.size() - 1, ridx[ridx.size()-1] + 1);
-      for (int i = static_cast<int>(ridx.size() - 1); (i > 0) &&
-                                  GetConstInt(ridx[i] >= reduceshape[i]); --i) {
+      ridx.Set(ridx.size() - 1, ridx[ridx.size() - 1] + 1);
+      for (int i = static_cast<int>(ridx.size() - 1);
+           (i > 0) && GetConstInt(ridx[i] >= reduceshape[i]); --i) {
         ridx.Set(i, ridx[i] - reduceshape[i]);
-        ridx.Set(i-1, ridx[i-1] + 1);
+        ridx.Set(i - 1, ridx[i - 1] + 1);
       }
       rec_flag = GetConstInt(ridx[0] < reduceshape[0]);
     } while (rec_flag);
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index 0d7fc9d9f010..261fdf9970a3 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -38,7 +38,6 @@
 #include <string>
 #include <unordered_set>
 #include <vector>
-#include <bitset>
 
 #include "detail/broadcast.h"
 
diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 1bea6696c13a..80f5d31af943 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -541,7 +541,7 @@ def tensordot(a, b, axes):
 def einsum(subscripts, a_tuple):
     # Grab non-einsum kwargs; do not optimize by default.
     # Temporally not implement the optimized einsum
-    if(not isinstance(a_tuple, tuple)):
+    if not isinstance(a_tuple, tuple):
         a_tuple = (a_tuple,)
     return cpp.einsum(subscripts, a_tuple)
 
diff --git a/src/topi/transform.cc b/src/topi/transform.cc
index 17624ed0a0c6..f71fae3c5aaa 100644
--- a/src/topi/transform.cc
+++ b/src/topi/transform.cc
@@ -23,8 +23,8 @@
  */
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/registry.h>
-#include <tvm/topi/transform.h>
 #include <tvm/topi/einsum.h>
+#include <tvm/topi/transform.h>
 #include <tvm/topi/utils.h>
 
 namespace tvm {
diff --git a/tests/python/unittest/test_te_autodiff.py b/tests/python/unittest/test_te_autodiff.py
index 6031182091fe..bae5d266f4c9 100644
--- a/tests/python/unittest/test_te_autodiff.py
+++ b/tests/python/unittest/test_te_autodiff.py
@@ -171,6 +171,10 @@ def fidentity(t0):
     check_grad(Y, X)
 
 
+    X = te.placeholder((3,3), name='X')
+    Y = topi.einsum("ii->i", (X))
+    check_grad(Y, X)
+
 def test_topi():
     X = te.placeholder((1, 2, 4, 4), name="X")
     W = te.placeholder((5, 2, 3, 3), name="W")

From 06c5f74f9698014274be9f90b8f701e391f9ee21 Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Thu, 17 Sep 2020 14:17:43 +0800
Subject: [PATCH 12/14] * create test einsum

---
 tests/python/topi/python/test_topi_einsum.py | 74 ++++++++++++++++++++
 tests/python/topi/python/test_topi_matmul.py |  1 -
 2 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 tests/python/topi/python/test_topi_einsum.py

diff --git a/tests/python/topi/python/test_topi_einsum.py b/tests/python/topi/python/test_topi_einsum.py
new file mode 100644
index 000000000000..306c6478654f
--- /dev/null
+++ b/tests/python/topi/python/test_topi_einsum.py
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+import numpy as np
+import tvm
+import tvm.testing
+from tvm import te
+from tvm import topi
+from tvm.topi.util import get_const_tuple
+
+def with_tvm(lam, *args):
+    """ Take numpy arrays as args, convert them to TVM tensors and call `lam`.
+    Result of lambda is converted back to numpy array and returned.
+    """
+    ctx = tvm.cpu(0)
+    pls = []     # placeholders
+    vals_nd = [] # initial values
+    for i,arg in enumerate(args):
+        pls.append(te.placeholder(arg.shape, name='pl'+str(i)))
+        vals_nd.append(tvm.nd.array(arg, ctx))
+
+    out = lam(*pls)
+    out_nd = tvm.nd.array(np.zeros(get_const_tuple(out.shape), dtype=out.dtype), ctx)
+    s = te.create_schedule([out.op])
+    m = tvm.build(s, pls + [out], "llvm")
+    m(*(vals_nd+[out_nd]))
+    return out_nd.asnumpy()
+
+def verify_einsum(subscripts, shapes):
+    ops = []
+    for shape in shapes :
+        tmp = np.random.uniform(low=-1.0, high=1.0, size=shape).astype(np.float32)
+        ops.append(tmp)
+
+    c1 = np.einsum(subscripts, *ops)
+
+    if len(ops) == 1 :
+        c2 = with_tvm(lambda A: topi.einsum(subscripts, (A)), *ops)
+    elif len(ops) == 2:
+        c2 = with_tvm(lambda A, B: topi.einsum(subscripts, (A, B)), *ops)
+    elif len(ops) == 3:
+        c2 = with_tvm(lambda A, B, C: topi.einsum(subscripts, (A, B, C)), *ops)
+
+    tvm.testing.assert_allclose(c1, c2, rtol=1e-5, atol=1e-5)
+
+def test_einsum():
+    verify_einsum('ii', [(5, 5)])
+    verify_einsum('ii->i', [(5, 5)])
+    verify_einsum('ij->i', [(5, 5)])
+    verify_einsum('...j->...', [(5, 5)])
+    verify_einsum('...j, j', [(5, 5), (5,)])
+    verify_einsum('..., ...', [(), (2, 3)])
+    verify_einsum('ijk, jil->kl', [(3, 4, 5), (4, 3, 2)])
+    verify_einsum('ij, ij -> i', [(1, 4), (2, 4)])
+    verify_einsum('...ij, ...jk -> ...ik', [(1, 4), (4, 2)])
+    verify_einsum('...ij, ...ik -> ...jk', [(1, 1, 1, 4), (1, 1, 1, 3)])
+    verify_einsum('ij,jk->ik', [(2,3), (3,4)])
+    verify_einsum('ij,jk,km->im', [(2,3),(3,4),(4,5)])
+
+if __name__ == "__main__":
+    test_einsum()
diff --git a/tests/python/topi/python/test_topi_matmul.py b/tests/python/topi/python/test_topi_matmul.py
index 26ba6f8142b4..997b0a06ec4c 100644
--- a/tests/python/topi/python/test_topi_matmul.py
+++ b/tests/python/topi/python/test_topi_matmul.py
@@ -77,7 +77,6 @@ def test_tensordot():
     verify_tensordot((3, 2, 2), (2, 3, 5), ((1, 0), (0, 1)))
     verify_tensordot((4, 3, 2, 2), (2, 4, 3, 5), ((1, 2, 0), (2, 0, 1)))
 
-
 if __name__ == "__main__":
     test_matmul()
     test_tensordot()

From 85786782460a4ca336231286f18c0a6fa2224f67 Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Tue, 22 Sep 2020 15:59:29 +0800
Subject: [PATCH 13/14] * Fix lint

---
 tests/python/topi/python/test_topi_einsum.py | 44 +++++++++++---------
 tests/python/topi/python/test_topi_matmul.py |  1 +
 tests/python/unittest/test_te_autodiff.py    |  4 +-
 3 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/tests/python/topi/python/test_topi_einsum.py b/tests/python/topi/python/test_topi_einsum.py
index 306c6478654f..4cce3a3f3eec 100644
--- a/tests/python/topi/python/test_topi_einsum.py
+++ b/tests/python/topi/python/test_topi_einsum.py
@@ -21,33 +21,35 @@
 from tvm import topi
 from tvm.topi.util import get_const_tuple
 
+
 def with_tvm(lam, *args):
-    """ Take numpy arrays as args, convert them to TVM tensors and call `lam`.
+    """Take numpy arrays as args, convert them to TVM tensors and call `lam`.
     Result of lambda is converted back to numpy array and returned.
     """
     ctx = tvm.cpu(0)
-    pls = []     # placeholders
-    vals_nd = [] # initial values
-    for i,arg in enumerate(args):
-        pls.append(te.placeholder(arg.shape, name='pl'+str(i)))
+    pls = []  # placeholders
+    vals_nd = []  # initial values
+    for i, arg in enumerate(args):
+        pls.append(te.placeholder(arg.shape, name="pl" + str(i)))
         vals_nd.append(tvm.nd.array(arg, ctx))
 
     out = lam(*pls)
     out_nd = tvm.nd.array(np.zeros(get_const_tuple(out.shape), dtype=out.dtype), ctx)
     s = te.create_schedule([out.op])
     m = tvm.build(s, pls + [out], "llvm")
-    m(*(vals_nd+[out_nd]))
+    m(*(vals_nd + [out_nd]))
     return out_nd.asnumpy()
 
+
 def verify_einsum(subscripts, shapes):
     ops = []
-    for shape in shapes :
+    for shape in shapes:
         tmp = np.random.uniform(low=-1.0, high=1.0, size=shape).astype(np.float32)
         ops.append(tmp)
 
     c1 = np.einsum(subscripts, *ops)
 
-    if len(ops) == 1 :
+    if len(ops) == 1:
         c2 = with_tvm(lambda A: topi.einsum(subscripts, (A)), *ops)
     elif len(ops) == 2:
         c2 = with_tvm(lambda A, B: topi.einsum(subscripts, (A, B)), *ops)
@@ -56,19 +58,21 @@ def verify_einsum(subscripts, shapes):
 
     tvm.testing.assert_allclose(c1, c2, rtol=1e-5, atol=1e-5)
 
+
 def test_einsum():
-    verify_einsum('ii', [(5, 5)])
-    verify_einsum('ii->i', [(5, 5)])
-    verify_einsum('ij->i', [(5, 5)])
-    verify_einsum('...j->...', [(5, 5)])
-    verify_einsum('...j, j', [(5, 5), (5,)])
-    verify_einsum('..., ...', [(), (2, 3)])
-    verify_einsum('ijk, jil->kl', [(3, 4, 5), (4, 3, 2)])
-    verify_einsum('ij, ij -> i', [(1, 4), (2, 4)])
-    verify_einsum('...ij, ...jk -> ...ik', [(1, 4), (4, 2)])
-    verify_einsum('...ij, ...ik -> ...jk', [(1, 1, 1, 4), (1, 1, 1, 3)])
-    verify_einsum('ij,jk->ik', [(2,3), (3,4)])
-    verify_einsum('ij,jk,km->im', [(2,3),(3,4),(4,5)])
+    verify_einsum("ii", [(5, 5)])
+    verify_einsum("ii->i", [(5, 5)])
+    verify_einsum("ij->i", [(5, 5)])
+    verify_einsum("...j->...", [(5, 5)])
+    verify_einsum("...j, j", [(5, 5), (5,)])
+    verify_einsum("..., ...", [(), (2, 3)])
+    verify_einsum("ijk, jil->kl", [(3, 4, 5), (4, 3, 2)])
+    verify_einsum("ij, ij -> i", [(1, 4), (2, 4)])
+    verify_einsum("...ij, ...jk -> ...ik", [(1, 4), (4, 2)])
+    verify_einsum("...ij, ...ik -> ...jk", [(1, 1, 1, 4), (1, 1, 1, 3)])
+    verify_einsum("ij,jk->ik", [(2, 3), (3, 4)])
+    verify_einsum("ij,jk,km->im", [(2, 3), (3, 4), (4, 5)])
+
 
 if __name__ == "__main__":
     test_einsum()
diff --git a/tests/python/topi/python/test_topi_matmul.py b/tests/python/topi/python/test_topi_matmul.py
index 997b0a06ec4c..26ba6f8142b4 100644
--- a/tests/python/topi/python/test_topi_matmul.py
+++ b/tests/python/topi/python/test_topi_matmul.py
@@ -77,6 +77,7 @@ def test_tensordot():
     verify_tensordot((3, 2, 2), (2, 3, 5), ((1, 0), (0, 1)))
     verify_tensordot((4, 3, 2, 2), (2, 4, 3, 5), ((1, 2, 0), (2, 0, 1)))
 
+
 if __name__ == "__main__":
     test_matmul()
     test_tensordot()
diff --git a/tests/python/unittest/test_te_autodiff.py b/tests/python/unittest/test_te_autodiff.py
index bae5d266f4c9..b2f26471d267 100644
--- a/tests/python/unittest/test_te_autodiff.py
+++ b/tests/python/unittest/test_te_autodiff.py
@@ -170,11 +170,11 @@ def fidentity(t0):
     Y = topi.tensordot(A, B, 1)
     check_grad(Y, X)
 
-
-    X = te.placeholder((3,3), name='X')
+    X = te.placeholder((3, 3), name="X")
     Y = topi.einsum("ii->i", (X))
     check_grad(Y, X)
 
+
 def test_topi():
     X = te.placeholder((1, 2, 4, 4), name="X")
     W = te.placeholder((5, 2, 3, 3), name="W")

From 605c4db8fc62fe3f6e1caa4a55fb420bfa6803ec Mon Sep 17 00:00:00 2001
From: Ke Han <kehan5800@gmail.com>
Date: Mon, 11 Jan 2021 23:51:24 +0800
Subject: [PATCH 14/14] * Fix comments

---
 include/tvm/topi/einsum.h                    | 95 +++++++++++---------
 include/tvm/topi/tags.h                      |  1 +
 python/tvm/topi/__init__.py                  |  1 +
 python/tvm/topi/einsum.py                    | 44 +++++++++
 python/tvm/topi/transform.py                 |  8 --
 tests/python/topi/python/test_topi_einsum.py |  8 +-
 6 files changed, 104 insertions(+), 53 deletions(-)
 create mode 100644 python/tvm/topi/einsum.py

diff --git a/include/tvm/topi/einsum.h b/include/tvm/topi/einsum.h
index 29df13460a7f..e1baadab09d3 100644
--- a/include/tvm/topi/einsum.h
+++ b/include/tvm/topi/einsum.h
@@ -24,6 +24,10 @@
 #ifndef TVM_TOPI_EINSUM_H_
 #define TVM_TOPI_EINSUM_H_
 
+#define LABELRANGE 128
+#define NPY_MAXDIMS 16
+#define NPY_MAXARGS 16
+
 #include <tvm/te/operation.h>
 #include <tvm/tir/data_layout.h>
 #include <tvm/topi/detail/constant_utils.h>
@@ -35,6 +39,7 @@
 #include <bitset>
 #include <iterator>
 #include <string>
+#include <tuple>
 #include <unordered_set>
 #include <vector>
 
@@ -220,8 +225,8 @@ inline int ParseOutputSubscripts(const char* subscripts, int length, int ndim_br
           << "in an input";
 
       /* Check that there is room in out_labels for this label. */
-      CHECK(ndim < 16) << "einstein sum subscripts string contains "
-                       << "too many subscripts in the output";
+      CHECK(ndim < NPY_MAXDIMS) << "einstein sum subscripts string contains "
+                                << "too many subscripts in the output";
 
       out_labels[ndim++] = label;
     } else if (label == '.') {
@@ -233,8 +238,8 @@ inline int ParseOutputSubscripts(const char* subscripts, int length, int ndim_br
           << "an ellipsis ('...') in the output";
 
       /* Check there is room in out_labels for broadcast dims. */
-      CHECK(ndim + ndim_broadcast <= 16) << "einstein sum subscripts string contains "
-                                         << "too many subscripts in the output";
+      CHECK(ndim + ndim_broadcast <= NPY_MAXDIMS) << "einstein sum subscripts string contains "
+                                                  << "too many subscripts in the output";
 
       ellipsis = 1;
       for (bdim = 0; bdim < ndim_broadcast; ++bdim) {
@@ -273,7 +278,7 @@ inline int ParseOutputSubscripts(const char* subscripts, int length, int ndim_br
 inline void GetCombinedDimsView(const Tensor& op, int iop, char* labels, Array<PrimExpr>* newshape,
                                 Array<PrimExpr>* newstride) {
   int idim, ndim, icombine, combineoffset;
-  int icombinemap[16];
+  int icombinemap[NPY_MAXDIMS];
   int newdim;
 
   Array<PrimExpr> shape = op->shape;
@@ -402,8 +407,8 @@ inline int CountSubstring(const std::string& str, const std::string& sub) {
  *
  * \return bitset.
  */
-inline std::bitset<128> Str2Set(const std::string& str) {
-  std::bitset<128> ret;
+inline std::bitset<LABELRANGE> Str2Set(const std::string& str) {
+  std::bitset<LABELRANGE> ret;
   for (const char& c : str) {
     ret.set(static_cast<int>(c));
   }
@@ -438,10 +443,10 @@ inline std::vector<std::string> Split(const std::string& str, const std::string&
  * if no output, the vector[1] is NULL.
  * "ab, bc -> ac" => ["ab,bc", "ac"]
  */
-inline std::vector<std::string> ParseEinsumInput(std::string subscripts,
-                                                 const std::vector<Array<PrimExpr>>& operands) {
+inline std::tuple<std::string, std::string> ParseEinsumInput(
+    std::string subscripts, const std::vector<Array<PrimExpr>>& operands) {
   const std::string einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
-  std::bitset<128> einsum_symbols_set;
+  std::bitset<LABELRANGE> einsum_symbols_set;
   for (const char& c : einsum_symbols) {
     einsum_symbols_set.set(c);
   }
@@ -475,7 +480,7 @@ inline std::vector<std::string> ParseEinsumInput(std::string subscripts,
                        [](const char& c) { return c == '.' || c == ',' || c == '-' || c == '>'; }),
         used.end());
 
-    std::bitset<128> used_set = Str2Set(used);
+    std::bitset<LABELRANGE> used_set = Str2Set(used);
     std::string ellipse_inds = "";
     for (const char& c : einsum_symbols) {
       if (!used_set.test(static_cast<int>(c))) {
@@ -544,7 +549,7 @@ inline std::vector<std::string> ParseEinsumInput(std::string subscripts,
       subscripts += "->" + output_sub;
     } else {
       // Special care for outputless ellipses
-      std::bitset<128> out_ellipse_set = Str2Set(out_ellipse);
+      std::bitset<LABELRANGE> out_ellipse_set = Str2Set(out_ellipse);
       std::string tmp_subscripts = subscripts, output_subscript = "";
       size_t len_tmp_subscripts = tmp_subscripts.length();
       std::sort(tmp_subscripts.begin(), tmp_subscripts.end());
@@ -565,12 +570,14 @@ inline std::vector<std::string> ParseEinsumInput(std::string subscripts,
   }
 
   // Build output string if does not exist
-  std::vector<std::string> ret(2);
+  std::tuple<std::string, std::string> ret;
   if (subscripts.find("->") != std::string::npos) {
-    ret = Split(subscripts, "->");
+    std::vector<std::string> tmp(2);
+    tmp = Split(subscripts, "->");
+    ret = std::make_tuple(tmp[0], tmp[1]);
   } else {
-    ret[0] = subscripts;
-    ret[1] = "";
+    std::string first = subscripts;
+    std::string second = "";
     // Build output subscripts
     std::string tmp_subscripts = subscripts;
     size_t len_tmp_subscripts = tmp_subscripts.length();
@@ -583,20 +590,21 @@ inline std::vector<std::string> ParseEinsumInput(std::string subscripts,
       CHECK(einsum_symbols_set.test(c)) << "Character " << c << " is not a valid symbol.";
       if ((i == 0 || tmp_subscripts[i - 1] != c) &&
           (i == len_tmp_subscripts - 1 || tmp_subscripts[i + 1] != c)) {
-        ret[1].append(1, c);
+        second.append(1, c);
       }
     }
+    ret = std::make_tuple(first, second);
   }
 
   // Make sure output subscripts are in the input
-  std::bitset<128> input_subscripts_set = Str2Set(ret[0]);
-  for (const char& c : ret[1]) {
+  std::bitset<LABELRANGE> input_subscripts_set = Str2Set(std::get<0>(ret));
+  for (const char& c : std::get<1>(ret)) {
     CHECK(input_subscripts_set.test(c))
         << "Output character " << c << " did not appear in the input";
   }
 
   // Make sure number operands is equivalent to the number of terms
-  CHECK_EQ(std::count(ret[0].begin(), ret[0].end(), ',') + 1, operands.size())
+  CHECK_EQ(std::count(std::get<0>(ret).begin(), std::get<0>(ret).end(), ',') + 1, operands.size())
       << "Number of einsum subscripts must be equal to the "
       << "number of operands.";
 
@@ -613,14 +621,14 @@ inline std::vector<std::string> ParseEinsumInput(std::string subscripts,
 inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
                                         const std::vector<Array<PrimExpr>>& operands) {
   // Parsing
-  std::vector<std::string> parsed_subscripts = ParseEinsumInput(subscripts, operands);
+  std::tuple<std::string, std::string> parsed_subscripts = ParseEinsumInput(subscripts, operands);
 
   // Build a few useful list and sets
-  std::vector<std::string> input_list = Split(parsed_subscripts[0], ",");
+  std::vector<std::string> input_list = Split(std::get<0>(parsed_subscripts), ",");
   size_t isize = input_list.size();
 
   // Get length of each unique dimension and ensure all dimensions are correct
-  int dimension_dict[128];
+  int dimension_dict[LABELRANGE];
   memset(dimension_dict, -1, sizeof(dimension_dict));
   for (size_t i = 0; i < isize; ++i) {
     const std::string& term = input_list[i];
@@ -649,7 +657,7 @@ inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
   }
 
   // Get oshape
-  const std::string& output_str = parsed_subscripts[1];
+  const std::string& output_str = std::get<1>(parsed_subscripts);
   size_t odim = output_str.size();
   Array<PrimExpr> oshape(odim, -1);
   for (size_t i = 0; i < odim; ++i) {
@@ -662,22 +670,24 @@ inline Array<PrimExpr> NumpyEinsumShape(const std::string subscripts,
 /*!
  * \brief Evaluates the Einstein summation convention on the operands.
  *
- * \param subscripts_str Specifies the subscripts for summation as comma separated list of subscript
- * labels. \param inputs Arrays for the operation. \param name The name of the operation. \param tag
- * The tag to mark the operation.
+ * \param subscripts_str Specifies the subscripts for summation as comma separated list of
+ * subscript labels.
+ * \param inputs Arrays for the operation.
+ * \param name The name of the operation.
+ * \param tag The tag to mark the operation.
  *
  * \return The calculation based on the Einstein summation convention.
  */
 inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inputs,
-                     std::string name = "T_einsum", std::string tag = kMatMul) {
-  // able to compute op: trace, diag, sum, transpose, matmul, dot, inner, outer, multiply, tensordot
+                     std::string name = "T_einsum", std::string tag = kEinsum) {
   bool back = false;
   const char* subscripts = subscripts_str.data();
+  const char* head = subscripts;
   const int nop = inputs.size();
 
   /* Step 1: Parse the subscripts string into label_counts and op_labels */
-  int iop, idim, min_label = 127, max_label = 0;
-  char label_counts[128], op_labels[16][16];
+  int iop, idim, min_label = LABELRANGE - 1, max_label = 0;
+  char label_counts[LABELRANGE], op_labels[NPY_MAXARGS][NPY_MAXDIMS];
   memset(label_counts, 0, sizeof(label_counts));
   for (iop = 0; iop < nop; ++iop) {
     int length = static_cast<int>(strcspn(subscripts, ",-"));
@@ -688,12 +698,15 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     CHECK(!(iop < nop - 1 && subscripts[length] != ','))
         << "fewer operands provided to einstein sum function "
         << "than specified in the subscripts string";
-    ParseOperandSubscripts(subscripts, length, inputs[iop + back].ndim(), iop, op_labels[iop],
-                           label_counts, &min_label, &max_label);
+    CHECK_EQ(ParseOperandSubscripts(subscripts, length, inputs[iop + back].ndim(), iop,
+                                    op_labels[iop], label_counts, &min_label, &max_label),
+             0);
 
     /* Move subscripts to the start of the labels for the next op */
     subscripts += length;
+
     if (iop < nop - 1) {
+      CHECK_LT(subscripts - head, subscripts_str.length()) << "subscripts out of range";
       subscripts++;
     }
   }
@@ -724,7 +737,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
    * using each label that appeared once, in alphabetical order.
    */
   int label, ndim_output;
-  char output_labels[16];
+  char output_labels[NPY_MAXDIMS];
   if (subscripts[0] == '\0') {
     /* If no output was specified, always broadcast left, as usual. */
     for (ndim_output = 0; ndim_output < ndim_broadcast; ++ndim_output) {
@@ -732,8 +745,8 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
     }
     for (label = min_label; label <= max_label; ++label) {
       if (label_counts[label] == 1) {
-        CHECK(ndim_output < 16) << "einstein sum subscript string has too many "
-                                << "distinct labels";
+        CHECK(ndim_output < NPY_MAXDIMS) << "einstein sum subscript string has too many "
+                                         << "distinct labels";
         output_labels[ndim_output++] = label;
       }
     }
@@ -795,7 +808,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   int ndim_iter = ndim_output;
   for (label = min_label; label <= max_label; ++label) {
     if (label_counts[label] > 0 && memchr(output_labels, label, ndim_output) == nullptr) {
-      CHECK(ndim_iter < 16) << "too many subscripts in einsum";
+      CHECK(ndim_iter < NPY_MAXDIMS) << "too many subscripts in einsum";
       iter_labels[ndim_iter++] = label;
     }
   }
@@ -813,8 +826,8 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
   Array<PrimExpr> ostride_true = GetStride(oshape);
   Array<PrimExpr> reduceshape;
   std::vector<Array<PrimExpr>> remainshape(nop);
-  int op_axes_arrays[16][16];
-  int* op_axes[16];
+  int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
+  int* op_axes[NPY_MAXARGS];
   for (iop = 0; iop < nop; ++iop) {
     op_axes[iop] = op_axes_arrays[iop];
     CHECK_GE(PrepareOpAxes(opshape[iop].size(), iop, op_labels[iop], op_axes[iop], ndim_iter,
@@ -891,7 +904,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
       }
     }
     Array<PrimExpr> ridx = UnravelIndex(0, reduceshape);
-    // AType needs to fit the output data
+
     PrimExpr sum = 0;
     bool rec_flag = false;
     do {
@@ -899,7 +912,7 @@ inline Tensor einsum(const std::string& subscripts_str, const Array<Tensor> inpu
       for (int iop = 0; iop < nop; ++iop) {
         if (iop != -1) {
           PrimExpr k = 0;
-          // k = dot(input_indices, ostride) + dot(ridx, rstride)
+
           for (size_t i = 0; i < input_indices.size(); ++i) {
             k += input_indices[i] * ostride[iop][i];
           }
diff --git a/include/tvm/topi/tags.h b/include/tvm/topi/tags.h
index 3b748ca60ce5..c3641ae0de12 100644
--- a/include/tvm/topi/tags.h
+++ b/include/tvm/topi/tags.h
@@ -41,6 +41,7 @@ constexpr auto kDepthwiseConv2dNCHW = "depthwise_conv2d_nchw";
 constexpr auto kDepthwiseConv2dNHWC = "depthwise_conv2d_nhwc";
 constexpr auto kDepthwiseConv2dBackInputNHWC = "depthwise_conv2d_back_input_nhwc";
 constexpr auto kDepthwiseConv2dBackWeightNHWC = "depthwise_conv2d_back_weight_nhwc";
+constexpr auto kEinsum = "einsum";
 constexpr auto kGroupConv2d = "group_conv2d";
 
 inline bool is_broadcast(std::string tag) {
diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py
index 873901df62a5..6836f04b5ada 100644
--- a/python/tvm/topi/__init__.py
+++ b/python/tvm/topi/__init__.py
@@ -41,6 +41,7 @@
 from .scatter_add import *
 from .argwhere import *
 from .cumsum import *
+from .einsum import *
 from . import generic
 from . import nn
 from . import x86
diff --git a/python/tvm/topi/einsum.py b/python/tvm/topi/einsum.py
new file mode 100644
index 000000000000..f1f426ec8173
--- /dev/null
+++ b/python/tvm/topi/einsum.py
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+# pylint: disable=invalid-name,consider-using-enumerate,redefined-outer-name
+"""Einsum operator"""
+from . import cpp
+
+
+def einsum(subscripts, *operand):
+    """Evaluates the Einstein summation convention on the operands.
+
+    Parameters
+    ----------
+    subscripts : string
+        Specifies the subscripts for summation as comma separated list of subscript labels.
+        An implicit (classical Einstein summation) calculation is performed unless the
+        explicit indicator ‘->’ is included as well as subscript labels of the precise
+        output form.
+
+    a_tuple : tuple of tvm.te.Tensor
+        These are the Tensors for the operation.
+        The only difference of einsum between in tvm and numpy is it needs an extra brackets
+        for the tensors. For example, topi.einsum("ij, jk -> ik", (A, B)).
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        The calculation based on the Einstein summation convention.
+    """
+
+    return cpp.einsum(subscripts, operand)
diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 80f5d31af943..6ddbc73e4666 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -538,14 +538,6 @@ def tensordot(a, b, axes):
     return cpp.tensordot(a, b, axes[0], axes[1])
 
 
-def einsum(subscripts, a_tuple):
-    # Grab non-einsum kwargs; do not optimize by default.
-    # Temporally not implement the optimized einsum
-    if not isinstance(a_tuple, tuple):
-        a_tuple = (a_tuple,)
-    return cpp.einsum(subscripts, a_tuple)
-
-
 def arange(start, stop=None, step=1, dtype="float32"):
     """Creates a tensor with evenly spaced values within a given interval.
 
diff --git a/tests/python/topi/python/test_topi_einsum.py b/tests/python/topi/python/test_topi_einsum.py
index 4cce3a3f3eec..49e951398f40 100644
--- a/tests/python/topi/python/test_topi_einsum.py
+++ b/tests/python/topi/python/test_topi_einsum.py
@@ -19,7 +19,7 @@
 import tvm.testing
 from tvm import te
 from tvm import topi
-from tvm.topi.util import get_const_tuple
+from tvm.topi.utils import get_const_tuple
 
 
 def with_tvm(lam, *args):
@@ -50,11 +50,11 @@ def verify_einsum(subscripts, shapes):
     c1 = np.einsum(subscripts, *ops)
 
     if len(ops) == 1:
-        c2 = with_tvm(lambda A: topi.einsum(subscripts, (A)), *ops)
+        c2 = with_tvm(lambda A: topi.einsum(subscripts, A), *ops)
     elif len(ops) == 2:
-        c2 = with_tvm(lambda A, B: topi.einsum(subscripts, (A, B)), *ops)
+        c2 = with_tvm(lambda A, B: topi.einsum(subscripts, A, B), *ops)
     elif len(ops) == 3:
-        c2 = with_tvm(lambda A, B, C: topi.einsum(subscripts, (A, B, C)), *ops)
+        c2 = with_tvm(lambda A, B, C: topi.einsum(subscripts, A, B, C), *ops)
 
     tvm.testing.assert_allclose(c1, c2, rtol=1e-5, atol=1e-5)