| 44 | // ensure proper device placement. |
| 45 | template <typename T> |
| 46 | Status Concat(OpKernelContext* context, const gtl::ArraySlice<Tensor>& inputs, |
| 47 | Tensor* output) { |
| 48 | const int input_dims = inputs[0].dims(); |
| 49 | const TensorShape& input_shape = inputs[0].shape(); |
| 50 | |
| 51 | // Note that we reduce the concat of k-dimensional tensors into a two |
| 52 | // dimensional concat. Assuming the dimensions of any input tensor are |
| 53 | // {y0, y1,...,ym-1}, we flatten it to {1, y}, where y = Prod_i(yi). |
| 54 | std::vector<std::unique_ptr<typename TTypes<T, 2>::ConstMatrix>> inputs_flat; |
| 55 | inputs_flat.reserve(inputs.size()); |
| 56 | int64 output_dim0 = 0; |
| 57 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 58 | const Tensor& input = inputs[i]; |
| 59 | if (input.dims() != input_dims) { |
| 60 | return errors::InvalidArgument( |
| 61 | "Ranks of all input tensors should match: shape[0] = ", |
| 62 | input_shape.DebugString(), " vs. shape[", i, |
| 63 | "] = ", input.shape().DebugString()); |
| 64 | } |
| 65 | for (int j = 1; j < input_dims; ++j) { |
| 66 | if (input.dim_size(j) != input_shape.dim_size(j)) { |
| 67 | return errors::InvalidArgument( |
| 68 | "Dimensions of inputs should match: shape[0] = ", |
| 69 | input_shape.DebugString(), " vs. shape[", i, |
| 70 | "] = ", input.shape().DebugString()); |
| 71 | } |
| 72 | } |
| 73 | if (input.NumElements() > 0) { |
| 74 | inputs_flat.emplace_back(new typename TTypes<T, 2>::ConstMatrix( |
| 75 | input.shaped<T, 2>({1, input.NumElements()}))); |
| 76 | } |
| 77 | output_dim0 += input.dim_size(0); |
| 78 | } |
| 79 | |
| 80 | TensorShape output_shape(input_shape); |
| 81 | output_shape.set_dim(0, output_dim0); |
| 82 | TF_RETURN_IF_ERROR( |
| 83 | context->allocate_temp(DataTypeToEnum<T>::value, output_shape, output)); |
| 84 | if (output->NumElements() > 0) { |
| 85 | auto output_flat = output->shaped<T, 2>({1, output->NumElements()}); |
| 86 | #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ |
| 87 | (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) |
| 88 | if (std::is_same<Device, GPUDevice>::value) { |
| 89 | ConcatGPU<T>(context, inputs_flat, output, &output_flat); |
| 90 | return Status::OK(); |
| 91 | } |
| 92 | #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
| 93 | ConcatCPU<T>(context->device(), inputs_flat, &output_flat); |
| 94 | } |
| 95 | |
| 96 | return Status::OK(); |
| 97 | } |
| 98 | |
| 99 | // The Split*() functions split 'input' with element type T into 'sizes.size()' |
| 100 | // tensors along the zeroth dimension, with the ith split having zeroth- |