| 296 | } |
| 297 | |
| 298 | void Compute(OpKernelContext* c) override { |
| 299 | const TensorList* l = nullptr; |
| 300 | OP_REQUIRES_OK(c, GetInputList(c, 0, &l)); |
| 301 | OP_REQUIRES(c, element_dtype_ == l->element_dtype, |
| 302 | errors::InvalidArgument("Invalid data types; op elements ", |
| 303 | DataTypeString(element_dtype_), |
| 304 | " but list elements ", |
| 305 | DataTypeString(l->element_dtype))); |
| 306 | int32 index = c->input(1).scalar<int32>()(); |
| 307 | OP_REQUIRES(c, index < l->tensors().size(), |
| 308 | errors::InvalidArgument("Trying to access element ", index, |
| 309 | " in a list with ", l->tensors().size(), |
| 310 | " elements.")); |
| 311 | if (l->tensors()[index].dtype() != DT_INVALID) { |
| 312 | c->set_output(0, l->tensors()[index]); |
| 313 | } else { |
| 314 | PartialTensorShape partial_element_shape; |
| 315 | OP_REQUIRES_OK( |
| 316 | c, GetElementShapeFromInput(c, *l, 2, &partial_element_shape)); |
| 317 | TensorShape element_shape; |
| 318 | // If l->element_shape and the element_shape input are both not fully |
| 319 | // defined, try to infer the shape from other list elements. This requires |
| 320 | // that all initialized list elements have the same shape. |
| 321 | // NOTE(srbs): This might be a performance bottleneck since we are |
| 322 | // iterating over the entire list here. This is necessary for feature |
| 323 | // parity with TensorArray.read. TensorArray has a mode in which all |
| 324 | // elements are required to be of the same shape, TensorList does not. |
| 325 | // In that mode TensorArray sets the array's element_shape on the first |
| 326 | // write call. We could do something similar here if needed. |
| 327 | if (!partial_element_shape.IsFullyDefined()) { |
| 328 | for (const Tensor& t : l->tensors()) { |
| 329 | if (t.dtype() != DT_INVALID) { |
| 330 | PartialTensorShape tmp = partial_element_shape; |
| 331 | OP_REQUIRES_OK(c, tmp.MergeWith(t.shape(), &partial_element_shape)); |
| 332 | } |
| 333 | } |
| 334 | } |
| 335 | OP_REQUIRES( |
| 336 | c, partial_element_shape.AsTensorShape(&element_shape), |
| 337 | errors::InvalidArgument("Trying to read an uninitialized tensor but ", |
| 338 | "element_shape is not fully defined: ", |
| 339 | partial_element_shape.DebugString(), |
| 340 | " and no list element is set.")); |
| 341 | Tensor* result; |
| 342 | AllocatorAttributes attr; |
| 343 | if (element_dtype_ == DT_VARIANT) { |
| 344 | attr.set_on_host(true); |
| 345 | } |
| 346 | OP_REQUIRES_OK(c, c->allocate_output(0, element_shape, &result, attr)); |
| 347 | functor::SetZeroFunctor<Device, T>()(c->eigen_device<Device>(), |
| 348 | result->flat<T>()); |
| 349 | } |
| 350 | } |
| 351 | |
| 352 | private: |
| 353 | DataType element_dtype_; |
nothing calls this directly
no test coverage detected