1 /* Copyright 2016 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_GATHER_ND_OP_H_
17 #define TENSORFLOW_CORE_KERNELS_GATHER_ND_OP_H_
18 // Functor definition for GatherOp, must be compilable by nvcc.
19
20 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
21 #include "tensorflow/core/framework/bounds_check.h"
22 #include "tensorflow/core/framework/op_kernel.h"
23 #include "tensorflow/core/framework/tensor.h"
24 #include "tensorflow/core/lib/core/status.h"
25 #include "tensorflow/core/platform/types.h"
26 #include "tensorflow/core/util/util.h"
27
28 namespace tensorflow {
29
30 class OpKernelContext;
31 class Status;
32 class Tensor;
33
34 namespace functor {
35
36 template <typename Device, typename T, typename Index, int IXDIM>
37 struct GatherNdSlice {
38 // Performs a slice gather op on (Tparams, Tindices), writing to Tout.
39 // Returns an index to Tindices if the value at that index is out of range.
40 // Returns -1 if all values of Tindices are in range.
41 Index operator()(const Device& d, const Index slice_size,
42 typename TTypes<int32>::Scalar Tscratch,
43 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams,
44 typename TTypes<Index>::ConstMatrix Tindices,
45 typename TTypes<T>::Matrix Tout);
46 };
47
48 template <typename Device, typename T, typename Index>
DoGatherNd(OpKernelContext * c,const Tensor & params,const Tensor & indices,Tensor * out)49 Status DoGatherNd(OpKernelContext* c, const Tensor& params,
50 const Tensor& indices, Tensor* out) {
51 if (!TensorShapeUtils::IsVectorOrHigher(params.shape())) {
52 return errors::InvalidArgument("params must be at least a vector");
53 }
54 if (!TensorShapeUtils::IsVectorOrHigher(indices.shape())) {
55 return errors::InvalidArgument("indices must be at least a vector");
56 }
57 if (indices.dim_size(indices.dims() - 1) > params.dims()) {
58 return errors::InvalidArgument(
59 "index innermost dimension length must be <= params rank; saw: ",
60 indices.dim_size(indices.dims() - 1), " vs. ", params.dims());
61 }
62
63 const TensorShape& indices_shape(indices.shape());
64 const int64 indices_nd = indices_shape.dim_size(indices_shape.dims() - 1);
65
66 // Check that we have enough index space
67 int64 N_big = 1;
68 for (int i = 0; i < indices_shape.dims() - 1; ++i) {
69 N_big *= indices_shape.dim_size(i);
70 }
71 if (N_big > std::numeric_limits<int>::max()) {
72 return errors::InvalidArgument(
73 "indices has too many elements for int indexing: ", N_big, " > ",
74 std::numeric_limits<int>::max());
75 }
76 if (params.NumElements() > std::numeric_limits<Index>::max()) {
77 return errors::InvalidArgument("params.NumElements() too large for ",
78 DataTypeString(DataTypeToEnum<Index>::v()),
79 " indexing: ", params.NumElements(), " > ",
80 std::numeric_limits<Index>::max());
81 }
82
83 // The result shape is
84 // indices.shape[:-1] + params.shape[indices.shape[-1]:]
85 Index N_result = 1;
86 for (int i = 0; i < indices_shape.dims() - 1; ++i) {
87 N_result *= indices_shape.dim_size(i);
88 }
89
90 const TensorShape& params_shape(params.shape());
91 Index total_nd = params_shape.dims();
92
93 TensorShape result_shape(indices_shape);
94 result_shape.RemoveLastDims(1);
95
96 int64 slice_size_big = 1;
97 for (Index i = indices_nd; i < total_nd; ++i) {
98 slice_size_big *= params_shape.dim_size(i);
99 result_shape.AddDim(params_shape.dim_size(i));
100 }
101
102 if (slice_size_big > std::numeric_limits<Index>::max()) {
103 return errors::InvalidArgument(
104 "slice size is too large for indexing: ", slice_size_big, " > ",
105 std::numeric_limits<Index>::max());
106 }
107
108 const Index slice_size = static_cast<Index>(slice_size_big);
109
110 TF_RETURN_IF_ERROR(
111 c->allocate_temp(DataTypeToEnum<T>::value, result_shape, out));
112
113 if (N_result > 0) {
114 if (params_shape.num_elements() == 0) {
115 return errors::InvalidArgument(
116 "Requested more than 0 entries, but "
117 "params is empty. Params shape: ",
118 params_shape.DebugString());
119 }
120
121 auto indices_mat = indices.flat_inner_dims<Index>();
122
123 Index bad_i = -1;
124
125 // Request to copy slices / subtensors
126 // Make out a matrix with the slices the col size.
127 auto out_mat = out->shaped<T, 2>({N_result, slice_size});
128 Tensor scratch;
129 TF_RETURN_IF_ERROR(c->allocate_temp(DT_INT32, TensorShape(), &scratch));
130 auto scratch_scalar = scratch.scalar<int32>();
131
132 switch (indices_nd) {
133 #define PARAMS_CASE(IXDIM) \
134 case IXDIM: { \
135 functor::GatherNdSlice<Device, T, Index, IXDIM> func; \
136 auto params_flat = params.flat_outer_dims<T, IXDIM + 1>(); \
137 bad_i = func(c->eigen_device<Device>(), slice_size, scratch_scalar, \
138 params_flat, indices_mat, out_mat); \
139 } break
140 PARAMS_CASE(0);
141 PARAMS_CASE(1);
142 PARAMS_CASE(2);
143 PARAMS_CASE(3);
144 PARAMS_CASE(4);
145 PARAMS_CASE(5);
146 PARAMS_CASE(6);
147 PARAMS_CASE(7);
148 #undef PARAMS_CASE
149 default:
150 return errors::InvalidArgument(
151 "Only indices.shape[-1] values between 1 and 7 "
152 "are currently supported. Requested rank: ",
153 indices_nd);
154 }
155
156 // bad_i will only return >= 0 on CPUs right now.
157 if (bad_i >= 0) {
158 auto shape = indices.shape();
159 shape.RemoveLastDims(1);
160 return errors::InvalidArgument(
161 "indices", SliceDebugString(shape, bad_i), " = [",
162 str_util::Join(
163 gtl::ArraySlice<Index>(&indices_mat(bad_i, 0), indices_nd), ", "),
164 "] does not index into param shape ", params.shape().DebugString(),
165 ", node name: ", c->op_kernel().name());
166 }
167 }
168 return Status::OK();
169 }
170
171 } // namespace functor
172 } // namespace tensorflow
173
174 #endif // TENSORFLOW_CORE_KERNELS_GATHER_ND_OP_H_
175