

#pragma once

#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "ngraph/ngraph.hpp"
#include "paddle/fluid/operators/ngraph/ops/op_bridge.h"
#include "paddle/fluid/platform/ngraph_helper.h"

namespace paddle {
namespace operators {
namespace ngraphs {

void BuildConcatNode(
    const std::shared_ptr<framework::OperatorBase>& op,
    std::shared_ptr<
        std::unordered_map<std::string, std::shared_ptr<ngraph::Node>>>
        ngb_node_map) {
  std::vector<std::shared_ptr<ngraph::Node>> args;
  for (auto& var_name_item : op->Inputs()) {
    for (auto& var_name : var_name_item.second) {
      auto& node0 = ngb_node_map->at(var_name);
      args.push_back(node0);
    }
  }
  auto op_attrs = framework::AttrReader(op->Attrs());
  int axis = op_attrs.Get<int>("axis");
  if (axis < 0) {
    axis = axis + args[0]->get_shape().size();
  }
  auto out = std::make_shared<ngraph::op::Concat>(args, axis);
  platform::SetOutputNode(op, "Out", out, ngb_node_map);
}
}  // namespace ngraphs
}  // namespace operators
}  // namespace paddle

REGISTER_NG_OP(concat, BuildConcatNode);
