| 607 | TensorListSplit(OpKernelConstruction* c) : OpKernel(c) {} |
| 608 | |
| 609 | void Compute(OpKernelContext* c) override { |
| 610 | Tensor* output_tensor; |
| 611 | AllocatorAttributes attr; |
| 612 | attr.set_on_host(true); |
| 613 | OP_REQUIRES_OK(c, c->allocate_output(0, {}, &output_tensor, attr)); |
| 614 | PartialTensorShape element_shape; |
| 615 | OP_REQUIRES_OK(c, TensorShapeFromTensor(c->input(1), &element_shape)); |
| 616 | OP_REQUIRES(c, element_shape.unknown_rank() || element_shape.dims() >= 1, |
| 617 | errors::InvalidArgument( |
| 618 | "TensorListSplit requires element_shape to be at least of ", |
| 619 | "rank 1, but saw: ", element_shape.DebugString())); |
| 620 | TensorList output_list; |
| 621 | const Tensor& input_tensor = c->input(0); |
| 622 | output_list.element_dtype = input_tensor.dtype(); |
| 623 | OP_REQUIRES(c, TensorShapeUtils::IsVectorOrHigher(input_tensor.shape()), |
| 624 | errors::InvalidArgument( |
| 625 | "Tensor must be at least a vector, but saw shape: ", |
| 626 | input_tensor.shape().DebugString())); |
| 627 | TensorShape tensor_shape_without_first_dim(input_tensor.shape()); |
| 628 | tensor_shape_without_first_dim.RemoveDim(0); |
| 629 | PartialTensorShape element_shape_without_first_dim; |
| 630 | if (!element_shape.unknown_rank()) { |
| 631 | element_shape_without_first_dim = |
| 632 | PartialTensorShape(element_shape.dim_sizes()); |
| 633 | element_shape_without_first_dim.RemoveDim(0); |
| 634 | } |
| 635 | OP_REQUIRES(c, |
| 636 | element_shape_without_first_dim.IsCompatibleWith( |
| 637 | tensor_shape_without_first_dim), |
| 638 | errors::InvalidArgument( |
| 639 | "tensor shape ", input_tensor.shape().DebugString(), |
| 640 | " is not compatible with element_shape ", |
| 641 | element_shape.DebugString())); |
| 642 | output_list.element_shape = element_shape; |
| 643 | const Tensor& lengths = c->input(2); |
| 644 | OP_REQUIRES(c, TensorShapeUtils::IsVector(lengths.shape()), |
| 645 | errors::InvalidArgument( |
| 646 | "Expected lengths to be a vector, received shape: ", |
| 647 | lengths.shape().DebugString())); |
| 648 | output_list.tensors().reserve(lengths.shape().dim_size(0)); |
| 649 | int64 start = 0; |
| 650 | int64 end = 0; |
| 651 | for (int i = 0; i < lengths.shape().dim_size(0); ++i) { |
| 652 | int64 length = lengths.vec<int64>()(i); |
| 653 | OP_REQUIRES( |
| 654 | c, length >= 0, |
| 655 | errors::InvalidArgument("Invalid value in lengths: ", length)); |
| 656 | end = start + length; |
| 657 | OP_REQUIRES(c, end <= input_tensor.shape().dim_size(0), |
| 658 | errors::InvalidArgument("Attempting to slice [", start, ", ", |
| 659 | end, "] from tensor with length ", |
| 660 | input_tensor.shape().dim_size(0))); |
| 661 | Tensor tmp = input_tensor.Slice(start, end); |
| 662 | start = end; |
| 663 | // TODO(apassos) maybe not always align; but weird compiler bugs seem to |
| 664 | // prevent this. |
| 665 | Tensor aligned; |
| 666 | OP_REQUIRES_OK(c, c->allocate_temp(tmp.dtype(), tmp.shape(), &aligned)); |
nothing calls this directly
no test coverage detected