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 // See docs in ../ops/array_ops.cc. 17 18 #include "tensorflow/core/framework/bounds_check.h" 19 #include "tensorflow/core/framework/op_kernel.h" 20 #include "tensorflow/core/framework/register_types.h" 21 #include "tensorflow/core/framework/tensor.h" 22 #include "tensorflow/core/framework/variant.h" 23 #include "tensorflow/core/framework/variant_encode_decode.h" 24 #include "tensorflow/core/kernels/gather_functor.h" 25 #include "tensorflow/core/platform/mem.h" 26 #include "tensorflow/core/platform/types.h" 27 #include "tensorflow/core/util/util.h" 28 29 namespace tensorflow { 30 31 typedef Eigen::ThreadPoolDevice CPUDevice; 32 typedef Eigen::GpuDevice GPUDevice; 33 34 template <typename Device, typename T, typename Index> 35 class GatherOp : public OpKernel { 36 public: 37 // QUESTION: It'd be nice to support DT_INT16, DT_UINT8, 38 // etc. here for the type of the second input argument. Should 39 // we have the framework do some sort of integer promotion 40 // automatically, or should that be something that users have to 41 // do explicitly with a conversion operator in the graph? GatherOp(OpKernelConstruction * c)42 explicit GatherOp(OpKernelConstruction* c) : OpKernel(c) {} 43 Compute(OpKernelContext * c)44 void Compute(OpKernelContext* c) override { 45 const Tensor& params = c->input(0); 46 const Tensor& indices = c->input(1); 47 OP_REQUIRES( 48 c, TensorShapeUtils::IsVectorOrHigher(params.shape()), 49 errors::InvalidArgument("params must be at least 1 dimensional")); 50 51 // GatherV2 added an axis argument. For backwards compatibility with Gather, 52 // fall back to axis 0 if the op does not have an axis input. 53 int64 axis = 0; 54 if (c->num_inputs() == 3) { 55 const Tensor& axis_tensor = c->input(2); 56 OP_REQUIRES(c, TensorShapeUtils::IsScalar(axis_tensor.shape()), 57 errors::InvalidArgument("axis must be scalar")); 58 59 if (axis_tensor.dtype() == DT_INT32) { 60 axis = axis_tensor.scalar<int32>()(); 61 } else if (axis_tensor.dtype() == DT_INT64) { 62 axis = axis_tensor.scalar<int64>()(); 63 } else { 64 OP_REQUIRES(c, false, 65 errors::InvalidArgument("axis must be int32 or int64.")); 66 } 67 } 68 69 OP_REQUIRES( 70 c, axis >= -params.dims() && axis < params.dims(), 71 errors::InvalidArgument("Expected axis in the range [", -params.dims(), 72 ", ", params.dims(), "), but got ", axis)); 73 if (axis < 0) { 74 axis = params.dims() + axis; 75 } 76 77 // Check that we have enough index space 78 const int64 gather_dim_size = params.dim_size(axis); 79 const int64 N = indices.NumElements(); 80 OP_REQUIRES( 81 c, gather_dim_size <= std::numeric_limits<Index>::max(), 82 errors::InvalidArgument("params.shape[", axis, "] too large for ", 83 DataTypeString(DataTypeToEnum<Index>::v()), 84 " indexing: ", gather_dim_size, " > ", 85 std::numeric_limits<Index>::max())); 86 87 // The result shape is params.shape[0:axis] + indices.shape + 88 // params.shape[axis + 1:]. 89 TensorShape result_shape; 90 int64 outer_size = 1; 91 int64 inner_size = 1; 92 for (int i = 0; i < axis; i++) { 93 result_shape.AddDim(params.dim_size(i)); 94 outer_size *= params.dim_size(i); 95 } 96 result_shape.AppendShape(indices.shape()); 97 for (int i = axis + 1; i < params.dims(); i++) { 98 result_shape.AddDim(params.dim_size(i)); 99 inner_size *= params.dim_size(i); 100 } 101 102 Tensor* out = nullptr; 103 OP_REQUIRES_OK(c, c->allocate_output(0, result_shape, &out)); 104 if (N > 0 && outer_size > 0 && inner_size > 0) { 105 auto params_flat = 106 params.shaped<T, 3>({outer_size, gather_dim_size, inner_size}); 107 auto indices_flat = indices.flat<Index>(); 108 auto out_flat = out->shaped<T, 3>({outer_size, N, inner_size}); 109 110 functor::GatherFunctor<Device, T, Index> functor; 111 int64 bad_i = functor(c, params_flat, indices_flat, out_flat); 112 113 OP_REQUIRES( 114 c, bad_i < 0, 115 errors::InvalidArgument( 116 "indices", SliceDebugString(indices.shape(), bad_i), " = ", 117 indices_flat(bad_i), " is not in [0, ", gather_dim_size, ")")); 118 } 119 } 120 }; 121 122 #define REGISTER_GATHER_FULL(dev, type, index_type) \ 123 REGISTER_KERNEL_BUILDER(Name("Gather") \ 124 .Device(DEVICE_##dev) \ 125 .TypeConstraint<type>("Tparams") \ 126 .TypeConstraint<index_type>("Tindices"), \ 127 GatherOp<dev##Device, type, index_type>); \ 128 REGISTER_KERNEL_BUILDER(Name("GatherV2") \ 129 .Device(DEVICE_##dev) \ 130 .TypeConstraint<type>("Tparams") \ 131 .TypeConstraint<index_type>("Tindices") \ 132 .HostMemory("axis"), \ 133 GatherOp<dev##Device, type, index_type>) 134 135 #define REGISTER_GATHER_ALL_INDICES(dev, type) \ 136 REGISTER_GATHER_FULL(dev, type, int32); \ 137 REGISTER_GATHER_FULL(dev, type, int64) 138 139 #define REGISTER_GATHER_CPU(type) REGISTER_GATHER_ALL_INDICES(CPU, type) 140 141 // Registration of the CPU implementations. 142 TF_CALL_ALL_TYPES(REGISTER_GATHER_CPU); 143 TF_CALL_QUANTIZED_TYPES(REGISTER_GATHER_CPU); 144 TF_CALL_quint16(REGISTER_GATHER_CPU); 145 TF_CALL_qint16(REGISTER_GATHER_CPU); 146 TF_CALL_uint32(REGISTER_GATHER_CPU); 147 TF_CALL_uint64(REGISTER_GATHER_CPU); 148 149 #undef REGISTER_GATHER_CPU 150 151 #if GOOGLE_CUDA 152 153 // Registration of the GPU implementations. 154 #define REGISTER_GATHER_GPU(type) REGISTER_GATHER_ALL_INDICES(GPU, type) 155 156 TF_CALL_bool(REGISTER_GATHER_GPU); 157 TF_CALL_int32(REGISTER_GATHER_GPU); 158 TF_CALL_int64(REGISTER_GATHER_GPU); 159 TF_CALL_GPU_NUMBER_TYPES(REGISTER_GATHER_GPU); 160 TF_CALL_complex64(REGISTER_GATHER_GPU); 161 TF_CALL_complex128(REGISTER_GATHER_GPU); 162 163 #undef REGISTER_GATHER_GPU 164 165 #endif // GOOGLE_CUDA 166 167 #undef REGISTER_GATHER_ALL_INDICES 168 #undef REGISTER_GATHER_FULL 169 170 } // namespace tensorflow 171