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 // See docs in ../ops/string_ops.cc. 17 18 #include <string> 19 20 #include "tensorflow/core/framework/kernel_def_builder.h" 21 #include "tensorflow/core/framework/op_kernel.h" 22 #include "tensorflow/core/framework/tensor.h" 23 #include "tensorflow/core/framework/tensor_shape.h" 24 #include "tensorflow/core/lib/core/errors.h" 25 #include "tensorflow/core/lib/core/status.h" 26 #include "tensorflow/core/lib/core/stringpiece.h" 27 #include "tensorflow/core/lib/gtl/inlined_vector.h" 28 #include "tensorflow/core/lib/strings/str_util.h" 29 30 namespace tensorflow { 31 32 namespace { 33 34 const gtl::InlinedVector<int64, 8> GetStrides(const TensorShape& shape) { 35 gtl::InlinedVector<int64, 8> result(shape.dims()); 36 int64 product = 1; 37 for (int32 i = shape.dims() - 1; i >= 0; --i) { 38 result[i] = product; 39 product *= shape.dim_size(i); 40 } 41 return result; 42 } 43 44 // Given a linear index to a subset of dimensions, full shape, 45 // precomputed list of running products of the full shape, and list of 46 // dimensions in the subset, outputs the linear index to the full shape with 47 // nonspecified dimensions set to 0. Dimensions must be ordered from outer-most 48 // to inner-most with respect to the subset linear index. 49 inline int64 LinearSubIndexToFullIndex( 50 int64 output_index, const gtl::InlinedVector<int32, 8>& dim_list, 51 const TensorShape& input_shape, 52 const gtl::InlinedVector<int64, 8>& strides) { 53 int64 result = 0; 54 int64 quotient = output_index; 55 for (int32 i = dim_list.size() - 1; i >= 0; --i) { 56 int32 dim = dim_list[i]; 57 int64 dim_value = quotient % input_shape.dim_size(dim); 58 quotient = quotient / input_shape.dim_size(dim); 59 result += strides[dim] * dim_value; 60 } 61 return result; 62 } 63 64 // Computes the number of input elements reduced per output element. 65 int64 GetReductionIterSize(const gtl::InlinedVector<int32, 8>& reduced_indices, 66 const TensorShape& input_shape) { 67 int64 result = 1; 68 for (int32 reduce_dim : reduced_indices) { 69 result *= input_shape.dim_size(reduce_dim); 70 } 71 return result; 72 } 73 74 // Computes a list of all true reduced indices, accounting for negative 75 // indices. 76 gtl::InlinedVector<int32, 8> GetReducedIndices(const Tensor& reduction_indices, 77 int32 input_dims) { 78 const auto reduction_indices_flat = reduction_indices.flat<int32>(); 79 const int32 reduction_dims = reduction_indices_flat.size(); 80 81 gtl::InlinedVector<int32, 8> reduced_indices(reduction_dims); 82 for (int32 i = 0; i < reduction_dims; ++i) { 83 reduced_indices[i] = reduction_indices_flat(reduction_dims - i - 1); 84 reduced_indices[i] += reduced_indices[i] < 0 ? input_dims : 0; 85 } 86 87 return reduced_indices; 88 } 89 90 // Appends all unreduced dimensions to the given vector. 91 void MakeUnreducedIndices(gtl::InlinedVector<bool, 8> index_is_reduced, 92 int32 input_dims, 93 gtl::InlinedVector<int32, 8>* unreduced_indices) { 94 for (int32 index = 0; index < input_dims; ++index) { 95 if (!index_is_reduced[index]) unreduced_indices->push_back(index); 96 } 97 } 98 99 TensorShape GetOutputShape(gtl::InlinedVector<bool, 8> index_is_reduced, 100 const TensorShape& input_shape, bool keep_dims) { 101 TensorShape output_shape; 102 for (size_t index = 0; index < index_is_reduced.size(); ++index) { 103 if (index_is_reduced[index]) { 104 if (keep_dims) output_shape.AddDim(1); 105 } else { 106 output_shape.AddDim(input_shape.dim_size(index)); 107 } 108 } 109 return output_shape; 110 } 111 112 } // namespace 113 114 class ReduceJoinOp : public OpKernel { 115 public: 116 using OpKernel::OpKernel; 117 118 explicit ReduceJoinOp(OpKernelConstruction* ctx) : OpKernel(ctx) { 119 OP_REQUIRES_OK(ctx, ctx->GetAttr("keep_dims", &keep_dims_)); 120 OP_REQUIRES_OK(ctx, ctx->GetAttr("separator", &separator_)); 121 } 122 123 void Compute(OpKernelContext* context) override { 124 const Tensor& input = context->input(0); 125 const auto input_flat = input.flat<string>(); 126 const TensorShape& input_shape = input.shape(); 127 const int32 input_dims = input_shape.dims(); 128 129 const Tensor& reduction_indices = context->input(1); 130 const auto reduction_indices_flat = reduction_indices.flat<int32>(); 131 const int32 reduction_dims = reduction_indices_flat.size(); 132 133 gtl::InlinedVector<bool, 8> index_is_reduced(input_dims, false); 134 for (int32 i = 0; i < reduction_dims; i++) { 135 int32 reduce_index = reduction_indices_flat(i); 136 const int32 true_reduce_index = 137 reduce_index < 0 ? reduce_index + input_dims : reduce_index; 138 OP_REQUIRES( 139 context, reduce_index >= -input_dims && reduce_index < input_dims, 140 errors::OutOfRange("Invalid reduction dimension ", reduce_index, 141 " for input with ", input_dims, " dimension(s)")); 142 OP_REQUIRES(context, !index_is_reduced[true_reduce_index], 143 errors::InvalidArgument("Duplicate reduction dimension ", 144 reduce_index)); 145 index_is_reduced[true_reduce_index] = true; 146 } 147 148 gtl::InlinedVector<int32, 8> reduced_indices = 149 GetReducedIndices(reduction_indices, input_dims); 150 gtl::InlinedVector<int32, 8> unreduced_indices; 151 MakeUnreducedIndices(index_is_reduced, input_dims, &unreduced_indices); 152 const auto strides = GetStrides(input_shape); 153 154 Tensor* output_tensor = nullptr; 155 TensorShape output_shape = 156 GetOutputShape(index_is_reduced, input_shape, keep_dims_); 157 OP_REQUIRES_OK(context, context->allocate_output("output", output_shape, 158 &output_tensor)); 159 auto output_flat = output_tensor->flat<string>(); 160 161 const int64 reduction_iter_size = 162 GetReductionIterSize(reduced_indices, input_shape); 163 gtl::InlinedVector<StringPiece, 8> curr_strings(reduction_iter_size); 164 for (int64 output_index = 0; output_index < output_shape.num_elements(); 165 ++output_index) { 166 int64 output_full_index = LinearSubIndexToFullIndex( 167 output_index, unreduced_indices, input_shape, strides); 168 for (int64 reduction_index = 0; reduction_index < reduction_iter_size; 169 ++reduction_index) { 170 int64 reduction_full_index = LinearSubIndexToFullIndex( 171 reduction_index, reduced_indices, input_shape, strides); 172 curr_strings[reduction_index] = 173 input_flat(output_full_index + reduction_full_index); 174 } 175 output_flat(output_index) = 176 str_util::Join(curr_strings, separator_.c_str()); 177 } 178 } 179 180 private: 181 bool keep_dims_; 182 string separator_; 183 }; 184 185 REGISTER_KERNEL_BUILDER(Name("ReduceJoin").Device(DEVICE_CPU), ReduceJoinOp); 186 187 } // namespace tensorflow 188
