Function EinsumGradKernel

paddle/phi/kernels/impl/einsum_grad_kernel_impl.h:130–282

template <typename T, typename Context>
void EinsumGradKernel(const Context& dev_ctx,
                      const std::vector<const DenseTensor*>& x,
                      const std::vector<const DenseTensor*>& inner_cache,
                      const DenseTensor& out_grad,
                      const std::string& equation,
                      std::vector<DenseTensor*> x_grad) {
  VLOG(5) << "Start EinsumGradKernel:";
  bool has_zero_size_tensor = out_grad.numel() == 0;
  for (auto& i : x_grad) {
    if (i != nullptr) {
      if (i->numel() == 0) {
        has_zero_size_tensor = true;
      }
      Full<T, Context>(dev_ctx, i->dims(), 0, i);
    }
  }
  if (has_zero_size_tensor) return;
  LabelMap labelshape(0);
  LabelMap labeltype(LabelType::Reduction);
  std::vector<LabelMap> label2perms(x.size(), LabelMap(-1));
  std::vector<char> all_labels;  // order: ABO, AO, BO, AB, Reduce
  std::vector<std::vector<int64_t>> broadcast_shapes(2);
  std::vector<int64_t> output_dims;

  std::vector<DDim> input_dims;
  for (auto& i : x) {
    input_dims.push_back(i->dims());
  }
  std::vector<std::string> input_strs;
  std::string right;
  ParseEinsumEquation(equation,
                      input_dims,
                      &labelshape,
                      &labeltype,
                      &all_labels,
                      &label2perms,
                      &broadcast_shapes,
                      &output_dims,
                      &right,
                      &input_strs);

  VLOG(4) << "After grad parse einsum equation.";

  auto gather_labels_except_reduction = [&labeltype](std::string all) {
    std::string res("");
    for (auto c : all)
      if (labeltype[static_cast<int>(c)] != LabelType::Reduction) res += c;
    auto tmp_unique = unique_labels(res);
    return std::string(tmp_unique.begin(), tmp_unique.end());
  };
  if (x.size() == 1) {  // Unary
    auto splits = paddle::string::split_string(equation, "->");
    auto left = splits[0];
    right = splits[1];
    auto new_equation = right + "->" + gather_labels_except_reduction(left);
    auto new_operands = std::vector<const DenseTensor*>();
    new_operands.push_back(&out_grad);
    DenseTensor before_tile;