| 462 | data_format_(TensorFormat::FORMAT_NCHW) {} |
| 463 | |
| 464 | void Compute(OpKernelContext* context) override { |
| 465 | try { |
| 466 | auto cpu_engine = engine(ENGINE_CPU, 0); |
| 467 | OpInputList input_tensors; |
| 468 | GetMklInputList(context, "values", &input_tensors); |
| 469 | const int N = input_tensors.size(); |
| 470 | // Get Tensor shapes. |
| 471 | std::vector<MklDnnShape> mkl_input_shapes(N); |
| 472 | GetMklShapeList(context, "values", &mkl_input_shapes); |
| 473 | |
| 474 | const Tensor& concat_dim_tensor = (AxisArgName == NAME_IS_CONCAT_DIM) |
| 475 | ? MklGetInput(context, 0) |
| 476 | : MklGetInput(context, N); |
| 477 | // Sanity checks |
| 478 | OP_REQUIRES( |
| 479 | context, TensorShapeUtils::IsScalar(concat_dim_tensor.shape()), |
| 480 | errors::InvalidArgument( |
| 481 | "Concat dim tensor should be a scalar integer, but got shape ", |
| 482 | concat_dim_tensor.shape().DebugString())); |
| 483 | int32 concat_dim = |
| 484 | internal::SubtleMustCopy(concat_dim_tensor.scalar<int32>()()); |
| 485 | |
| 486 | // check that ranks of all tensors match |
| 487 | // and that their shapes match except for concat_dim. |
| 488 | int i = 0; |
| 489 | int num_of_empty_inputs = 0; |
| 490 | bool invoke_eigen = false; |
| 491 | bool are_all_mkl_inputs = true, are_all_tf_inputs = true; |
| 492 | const TensorShape expected_shape = mkl_input_shapes[0].IsMklTensor() |
| 493 | ? mkl_input_shapes[0].GetTfShape() |
| 494 | : input_tensors[0].shape(); |
| 495 | size_t expected_dims = expected_shape.dims(); |
| 496 | |
| 497 | if (concat_dim < 0) concat_dim = expected_dims + concat_dim; |
| 498 | |
| 499 | for (auto& s : mkl_input_shapes) { |
| 500 | TensorShape s_shape = |
| 501 | s.IsMklTensor() ? s.GetTfShape() : input_tensors[i].shape(); |
| 502 | size_t s_dims = s_shape.dims(); |
| 503 | |
| 504 | OP_REQUIRES( |
| 505 | context, s_dims == expected_dims, |
| 506 | errors::InvalidArgument( |
| 507 | "_MklConcatOp : Ranks of all input tensors should match:" |
| 508 | " input dimensions = ", |
| 509 | s_dims, " vs. expected rank = ", expected_dims)); |
| 510 | |
| 511 | for (int d = 0; d < expected_dims; ++d) { |
| 512 | if (d == concat_dim) continue; |
| 513 | |
| 514 | size_t expected_size = expected_shape.dim_size(d); |
| 515 | size_t s_size = s_shape.dim_size(d); |
| 516 | OP_REQUIRES( |
| 517 | context, expected_size == s_size, |
| 518 | errors::InvalidArgument("_MklConcatOp : Dimensions of inputs " |
| 519 | "should match: shape[0][", |
| 520 | d, "]= ", expected_size, " vs. shape[", i, |
| 521 | "][", d, "] = ", s_size)); |
no test coverage detected