1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_KERNELS_RESHAPE_OP_H_ 17 #define TENSORFLOW_CORE_KERNELS_RESHAPE_OP_H_ 18 19 #include <memory> 20 21 #include "tensorflow/core/framework/op_kernel.h" 22 #include "tensorflow/core/framework/register_types.h" 23 #include "tensorflow/core/framework/tensor.h" 24 #include "tensorflow/core/framework/tensor_shape.h" 25 #include "tensorflow/core/framework/types.h" 26 #include "tensorflow/core/lib/core/status.h" 27 #include "tensorflow/core/platform/logging.h" 28 #include "tensorflow/core/util/overflow.h" 29 30 namespace tensorflow { 31 32 // Note that this op is subclassed for QuantizedReshapeOp. 33 class ReshapeOp : public OpKernel { 34 public: ReshapeOp(OpKernelConstruction * context)35 explicit ReshapeOp(OpKernelConstruction* context) : OpKernel(context) {} 36 Compute(OpKernelContext * context)37 void Compute(OpKernelContext* context) override { 38 const Tensor& input = context->input(0); 39 const Tensor& sizes = context->input(1); 40 // Preliminary validation of sizes. 41 OP_REQUIRES( 42 context, 43 (TensorShapeUtils::IsVector(sizes.shape()) || 44 // TODO(rmlarsen): Disallow legacy use of scalars to represent shape. 45 TensorShapeUtils::IsScalar(sizes.shape())), 46 errors::InvalidArgument("sizes input must be 1-D, not ", 47 sizes.shape().DebugString())); 48 49 // Compute the output shape. Determine product of specified 50 // dimensions, and find the index of the unspecified one. 51 TensorShape shape; 52 int64_t product = 1; 53 int unknown_index = -1; 54 bool sizes_has_zero_dim; 55 switch (sizes.dtype()) { 56 case DT_INT32: 57 OP_REQUIRES_OK(context, 58 ValidateSizes<int32>(sizes, &product, &unknown_index, 59 &shape, &sizes_has_zero_dim)); 60 break; 61 case DT_INT64: 62 OP_REQUIRES_OK(context, 63 ValidateSizes<int64>(sizes, &product, &unknown_index, 64 &shape, &sizes_has_zero_dim)); 65 break; 66 default: 67 context->CtxFailure(errors::InvalidArgument( 68 "desired shape must be a DT_INT32 or DT_INT64 vector, not a ", 69 DataTypeString(sizes.dtype()))); 70 return; 71 } 72 if (unknown_index != -1) { 73 int64_t input_num_elements = 1; 74 bool input_has_zero_dim = false; 75 for (int dim = 0; dim < input.dims(); dim++) { 76 // For zero dimension, we don't count it into `input_num_elements` 77 // unless `sizes` has no zero dimension, so we are still able to 78 // infer shapes for other dimensions. 79 if (input.dim_size(dim) > 0 || !sizes_has_zero_dim) { 80 input_num_elements *= input.dim_size(dim); 81 } else { 82 input_has_zero_dim = true; 83 } 84 } 85 86 const int64_t missing = input_num_elements / product; 87 if (!input_has_zero_dim) { 88 OP_REQUIRES( 89 context, product * missing == input_num_elements, 90 errors::InvalidArgument( 91 "Input to reshape is a tensor with ", input_num_elements, 92 " values, but the requested shape requires a multiple of ", 93 product)); 94 } 95 shape.set_dim(unknown_index, missing); 96 } 97 OP_REQUIRES(context, shape.num_elements() == input.NumElements(), 98 errors::InvalidArgument("Input to reshape is a tensor with ", 99 input.NumElements(), 100 " values, but the requested shape has ", 101 shape.num_elements())); 102 103 // Actually produce the reshaped output. 104 Tensor output(input.dtype()); 105 CHECK(output.CopyFrom(input, shape)); 106 context->set_output(0, output); 107 } 108 IsExpensive()109 bool IsExpensive() override { return false; } 110 111 private: 112 template <typename Tshape> ValidateSizes(const Tensor & sizes,int64 * product,int * unknown_index,TensorShape * shape,bool * has_zero_dim)113 Status ValidateSizes(const Tensor& sizes, int64* product, int* unknown_index, 114 TensorShape* shape, bool* has_zero_dim) { 115 *product = 1; 116 *unknown_index = -1; 117 *has_zero_dim = false; 118 const int64_t num_dims = sizes.NumElements(); 119 auto Svec = sizes.flat<Tshape>(); 120 for (int d = 0; d < num_dims; ++d) { 121 const Tshape size = Svec(d); 122 if (size == -1) { 123 if (*unknown_index != -1) { 124 return errors::InvalidArgument( 125 "Only one input size may be -1, not both ", *unknown_index, 126 " and ", d); 127 } 128 *unknown_index = d; 129 shape->AddDim(1); 130 } else if (size < 0) { 131 return errors::InvalidArgument("Size ", d, 132 " must be non-negative, not ", size); 133 } else if (size == 0) { 134 // We don't include zero-sized dimension in product, so that we can 135 // still calculate number of elements for non-zero-sized dimensions and 136 // therefore infer their shapes. 137 shape->AddDim(size); 138 *has_zero_dim = true; 139 } else { 140 if (MultiplyWithoutOverflow(shape->num_elements(), size) < 0) { 141 string msg; 142 for (int ii = 0; ii < num_dims; ++ii) { 143 if (ii != 0) { 144 strings::StrAppend(&msg, ", "); 145 } 146 strings::StrAppend(&msg, Svec(ii)); 147 } 148 return errors::InvalidArgument("Shape [", msg, 149 "] has too many elements"); 150 } 151 shape->AddDim(size); 152 (*product) *= size; 153 } 154 } 155 return Status::OK(); 156 } 157 }; 158 159 } // namespace tensorflow 160 161 #endif // TENSORFLOW_CORE_KERNELS_RESHAPE_OP_H_ 162