| 196 | } |
| 197 | |
| 198 | void Compute(OpKernelContext* c) override { |
| 199 | const TensorList* tensor_list = nullptr; |
| 200 | OP_REQUIRES_OK(c, GetInputList(c, 0, &tensor_list)); |
| 201 | OP_REQUIRES( |
| 202 | c, element_dtype_ == tensor_list->element_dtype, |
| 203 | errors::InvalidArgument( |
| 204 | "Invalid data types; op elements ", DataTypeString(element_dtype_), |
| 205 | " but list elements ", DataTypeString(tensor_list->element_dtype))); |
| 206 | if (num_elements_ != -1) { |
| 207 | OP_REQUIRES(c, tensor_list->tensors().size() == num_elements_, |
| 208 | errors::InvalidArgument( |
| 209 | "Operation expected a list with ", num_elements_, |
| 210 | " elements but got a list with ", |
| 211 | tensor_list->tensors().size(), " elements.")); |
| 212 | } |
| 213 | PartialTensorShape partial_element_shape; |
| 214 | OP_REQUIRES_OK(c, GetElementShapeFromInput(c, *tensor_list, 1, |
| 215 | &partial_element_shape)); |
| 216 | OP_REQUIRES( |
| 217 | c, |
| 218 | partial_element_shape.IsFullyDefined() || |
| 219 | !tensor_list->tensors().empty(), |
| 220 | errors::InvalidArgument("Tried to stack elements of an empty ", |
| 221 | "list with non-fully-defined element_shape: ", |
| 222 | partial_element_shape.DebugString())); |
| 223 | |
| 224 | // Check that `element_shape` input tensor is compatible with the shapes of |
| 225 | // element tensors. |
| 226 | if (!tensor_list->element_shape.IsFullyDefined()) { |
| 227 | for (int i = 0; i < tensor_list->tensors().size(); ++i) { |
| 228 | const Tensor& t = tensor_list->tensors()[i]; |
| 229 | if (t.dtype() != DT_INVALID) { |
| 230 | PartialTensorShape tmp = partial_element_shape; |
| 231 | OP_REQUIRES_OK(c, tmp.MergeWith(t.shape(), &partial_element_shape)); |
| 232 | } |
| 233 | } |
| 234 | } |
| 235 | |
| 236 | // Compute the shape of the output tensor by pre-pending the leading dim to |
| 237 | // the element_shape. |
| 238 | TensorShape element_shape; |
| 239 | OP_REQUIRES(c, partial_element_shape.AsTensorShape(&element_shape), |
| 240 | errors::InvalidArgument( |
| 241 | "Tried to stack list which only contains uninitialized ", |
| 242 | "tensors and has a non-fully-defined element_shape: ", |
| 243 | partial_element_shape.DebugString())); |
| 244 | TensorShape output_shape = element_shape; |
| 245 | output_shape.InsertDim(0, tensor_list->tensors().size()); |
| 246 | Tensor* output; |
| 247 | OP_REQUIRES_OK(c, c->allocate_output(0, output_shape, &output)); |
| 248 | if (output->NumElements() == 0) { |
| 249 | return; |
| 250 | } |
| 251 | |
| 252 | ConstMatrixVector inputs_flat; |
| 253 | inputs_flat.reserve(tensor_list->tensors().size()); |
| 254 | Tensor zeros; |
| 255 | for (const auto& t : tensor_list->tensors()) { |
nothing calls this directly
no test coverage detected