| 33 | explicit ReshapeOp(OpKernelConstruction* context) : OpKernel(context) {} |
| 34 | |
| 35 | void Compute(OpKernelContext* context) override { |
| 36 | const Tensor& input = context->input(0); |
| 37 | const Tensor& sizes = context->input(1); |
| 38 | // Preliminary validation of sizes. |
| 39 | OP_REQUIRES(context, IsLegacyVector(sizes.shape()), |
| 40 | errors::InvalidArgument("sizes input must be 1-D, not ", |
| 41 | sizes.shape().DebugString())); |
| 42 | |
| 43 | // Compute the output shape. Determine product of specified |
| 44 | // dimensions, and find the index of the unspecified one. |
| 45 | TensorShape shape; |
| 46 | int64 product = 1; |
| 47 | int unknown_index = -1; |
| 48 | bool sizes_has_zero_dim; |
| 49 | switch (sizes.dtype()) { |
| 50 | case DT_INT32: |
| 51 | OP_REQUIRES_OK(context, |
| 52 | ValidateSizes<int32>(sizes, &product, &unknown_index, |
| 53 | &shape, &sizes_has_zero_dim)); |
| 54 | break; |
| 55 | case DT_INT64: |
| 56 | OP_REQUIRES_OK(context, |
| 57 | ValidateSizes<int64>(sizes, &product, &unknown_index, |
| 58 | &shape, &sizes_has_zero_dim)); |
| 59 | break; |
| 60 | default: |
| 61 | context->CtxFailure(errors::InvalidArgument( |
| 62 | "desired shape must be a DT_INT32 or DT_INT64 vector, not a ", |
| 63 | DataTypeString(sizes.dtype()))); |
| 64 | return; |
| 65 | } |
| 66 | if (unknown_index != -1) { |
| 67 | int64 input_num_elements = 1; |
| 68 | bool input_has_zero_dim = false; |
| 69 | for (int dim = 0; dim < input.dims(); dim++) { |
| 70 | // For zero dimension, we don't count it into `input_num_elements` |
| 71 | // unless `sizes` has no zero dimension, so we are still able to |
| 72 | // infer shapes for other dimensions. |
| 73 | if (input.dim_size(dim) > 0 || !sizes_has_zero_dim) { |
| 74 | input_num_elements *= input.dim_size(dim); |
| 75 | } else { |
| 76 | input_has_zero_dim = true; |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | const int64 missing = input_num_elements / product; |
| 81 | if (!input_has_zero_dim) { |
| 82 | OP_REQUIRES( |
| 83 | context, product * missing == input_num_elements, |
| 84 | errors::InvalidArgument( |
| 85 | "Input to reshape is a tensor with ", input_num_elements, |
| 86 | " values, but the requested shape requires a multiple of ", |
| 87 | product)); |
| 88 | } |
| 89 | shape.set_dim(unknown_index, missing); |
| 90 | } |
| 91 | OP_REQUIRES(context, shape.num_elements() == input.NumElements(), |
| 92 | errors::InvalidArgument("Input to reshape is a tensor with ", |
nothing calls this directly
no test coverage detected