Function Concat

tensorflow/core/kernels/mkl_concat_op_test.cc:38–76

template <typename T>
static Graph* Concat(const string& kind, int num_inputs,
		const TensorShape& in_shape, int concat_dims) {
  Graph* g = new Graph(OpRegistry::Global());
  DataType type = DataTypeToEnum<T>::v();

  const bool isDefault = (kind == "Default");
  string op_name = isDefault ? "Concat" : "_MklConcat";

  Tensor concat_dim(DT_INT32, TensorShape({}));
  concat_dim.scalar<int32>()() = concat_dims;

  Node* not_mkl_shape =
      test::graph::Constant(g, GetMklMetaTensor(), "not_mkl");

  std::vector<NodeBuilder::NodeOut> inputs;
  std::vector<NodeBuilder::NodeOut> inputs_not_mkl;
  inputs.reserve(num_inputs);
  inputs_not_mkl.reserve(num_inputs);

  for (int i = 0; i < num_inputs; ++i) {
    Tensor in(type, in_shape);
    in.flat<T>().setRandom();
    inputs.push_back(test::graph::Constant(g, in));
    inputs_not_mkl.push_back(test::graph::Constant(g, GetMklMetaTensor(), "not_mkl"));
  }

  auto nodeBuilder = NodeBuilder(g->NewName("n"), op_name)
                    .Input(test::graph::Constant(g, concat_dim))
                    .Input(inputs)
                    .Attr("N", num_inputs)
                    .Attr("T", type);

  isDefault ? nodeBuilder : nodeBuilder.Input(not_mkl_shape)
	                               .Input(inputs_not_mkl)
				       .Attr("_kernel", "MklLayoutDependentOp");
  TF_CHECK_OK(nodeBuilder.Finalize(g, nullptr));

  return g;
}

#define S_TENSOR(...) test::AsTensor<int32>({__VA_ARGS__})