| 56 | } |
| 57 | |
| 58 | Status Concat(const gtl::ArraySlice<Tensor>& tensors, Tensor* result) { |
| 59 | if (tensors.empty()) { |
| 60 | return errors::InvalidArgument("Cannot concatenate zero tensors"); |
| 61 | } |
| 62 | int64 total_dim0_size = 0; |
| 63 | for (const Tensor& tensor : tensors) { |
| 64 | if (tensor.dims() == 0) { |
| 65 | return errors::InvalidArgument( |
| 66 | "Cannot concatenate a zero-dimensional tensor"); |
| 67 | } |
| 68 | total_dim0_size += tensor.dim_size(0); |
| 69 | } |
| 70 | TensorShape shape = tensors[0].shape(); |
| 71 | shape.set_dim(0, total_dim0_size); |
| 72 | |
| 73 | const DataType dtype = tensors[0].dtype(); |
| 74 | for (int i = 1; i < tensors.size(); ++i) { |
| 75 | if (tensors[i].dtype() != dtype) { |
| 76 | return errors::InvalidArgument( |
| 77 | "Cannot concatenate tensors that have different data types"); |
| 78 | } |
| 79 | } |
| 80 | *result = Tensor(dtype, shape); |
| 81 | |
| 82 | // We use StringPiece as a convenient map over the tensor buffer, |
| 83 | // but we cast the type to get to the underlying buffer to do the |
| 84 | // copy. |
| 85 | StringPiece to_data = result->tensor_data(); |
| 86 | |
| 87 | if (DataTypeCanUseMemcpy(dtype)) { |
| 88 | int64 offset = 0; |
| 89 | for (const Tensor& tensor : tensors) { |
| 90 | StringPiece from_data = tensor.tensor_data(); |
| 91 | CHECK_LE(offset + from_data.size(), to_data.size()); |
| 92 | memcpy(const_cast<char*>(to_data.data()) + offset, from_data.data(), |
| 93 | from_data.size()); |
| 94 | |
| 95 | offset += from_data.size(); |
| 96 | } |
| 97 | } else { |
| 98 | if (dtype != DT_STRING) { |
| 99 | return errors::Internal("Unexpected data type"); |
| 100 | } |
| 101 | tstring* to_strings = |
| 102 | reinterpret_cast<tstring*>(const_cast<char*>(to_data.data())); |
| 103 | |
| 104 | int64 offset = 0; |
| 105 | for (const Tensor& tensor : tensors) { |
| 106 | auto from_strings = tensor.flat<tstring>(); |
| 107 | CHECK_LE(offset + tensor.NumElements(), result->NumElements()); |
| 108 | for (int i = 0; i < tensor.NumElements(); ++i) { |
| 109 | to_strings[offset + i] = from_strings(i); |
| 110 | } |
| 111 | |
| 112 | offset += tensor.NumElements(); |
| 113 | } |
| 114 | } |
| 115 | |