| 60 | |
| 61 | public: |
| 62 | void Compute(OpKernelContext* context) override { |
| 63 | const Tensor& input_tensor = MklGetInput(context, 0); |
| 64 | const Tensor& sizes = MklGetInput(context, 1); |
| 65 | |
| 66 | MklDnnShape mkl_shape_input; |
| 67 | GetMklShape(context, kInputSlotIdx, &mkl_shape_input); |
| 68 | bool input_in_mkl_format = mkl_shape_input.IsMklTensor(); |
| 69 | TensorShape input_shape = input_in_mkl_format ? mkl_shape_input.GetTfShape() |
| 70 | : input_tensor.shape(); |
| 71 | const int64 nelems = input_in_mkl_format ? input_shape.num_elements() |
| 72 | : input_tensor.NumElements(); |
| 73 | |
| 74 | // Preliminary validation of sizes. |
| 75 | OP_REQUIRES(context, TensorShapeUtils::IsVector(sizes.shape()), |
| 76 | errors::InvalidArgument("sizes input must be 1-D, not shape ", |
| 77 | sizes.shape().DebugString())); |
| 78 | |
| 79 | // Compute the output shape. Determine product of specified |
| 80 | // dimensions, and find the index of the unspecified one. |
| 81 | TensorShape shape; |
| 82 | int64 product = 1; |
| 83 | int unknown_index = -1; |
| 84 | bool sizes_has_zero_dim = false; |
| 85 | switch (sizes.dtype()) { |
| 86 | case DT_INT32: |
| 87 | OP_REQUIRES_OK(context, |
| 88 | ValidateSizes<int32>(sizes, &product, &unknown_index, |
| 89 | &shape, &sizes_has_zero_dim)); |
| 90 | break; |
| 91 | case DT_INT64: |
| 92 | OP_REQUIRES_OK(context, |
| 93 | ValidateSizes<int64>(sizes, &product, &unknown_index, |
| 94 | &shape, &sizes_has_zero_dim)); |
| 95 | break; |
| 96 | default: |
| 97 | context->CtxFailure(errors::InvalidArgument( |
| 98 | "desired shape must be a DT_INT32 or DT_INT64 vector, not a ", |
| 99 | DataTypeString(sizes.dtype()))); |
| 100 | return; |
| 101 | } |
| 102 | if (unknown_index != -1) { |
| 103 | int64 input_num_elements = 1; |
| 104 | bool input_has_zero_dim = false; |
| 105 | for (int dim = 0; dim < input_shape.dims(); ++dim) { |
| 106 | // For zero dimension, we don't count it into `input_num_elements` |
| 107 | // unless `sizes` has no zero dimension, so we are still able to |
| 108 | // infer shapes for other dimensions. |
| 109 | if (input_shape.dim_size(dim) > 0 || !sizes_has_zero_dim) { |
| 110 | input_num_elements *= input_shape.dim_size(dim); |
| 111 | } else { |
| 112 | input_has_zero_dim = true; |
| 113 | } |
| 114 | } |
| 115 | |
| 116 | const int64 missing = input_num_elements / product; |
| 117 | if (!input_has_zero_dim) { |
| 118 | OP_REQUIRES( |
| 119 | context, product * missing == input_num_elements, |
nothing calls this directly
no test coverage detected