| 902 | TensorListScatter(OpKernelConstruction* c) : OpKernel(c) {} |
| 903 | |
| 904 | void Compute(OpKernelContext* c) override { |
| 905 | Tensor* output_tensor; |
| 906 | AllocatorAttributes attr; |
| 907 | attr.set_on_host(true); |
| 908 | OP_REQUIRES_OK(c, c->allocate_output(0, {}, &output_tensor, attr)); |
| 909 | Tensor indices = c->input(1); |
| 910 | PartialTensorShape element_shape; |
| 911 | OP_REQUIRES_OK(c, TensorShapeFromTensor(c->input(2), &element_shape)); |
| 912 | // TensorListScatterV2 passes the num_elements input, TensorListScatter does |
| 913 | // not. |
| 914 | int num_elements = c->num_inputs() >= 4 ? c->input(3).scalar<int>()() : -1; |
| 915 | OP_REQUIRES(c, num_elements >= -1, |
| 916 | errors::InvalidArgument( |
| 917 | "TensorListScatter expects num_elements >= -1, found: ", |
| 918 | num_elements)); |
| 919 | TensorList output_list; |
| 920 | const Tensor& input_tensor = c->input(0); |
| 921 | output_list.element_dtype = input_tensor.dtype(); |
| 922 | OP_REQUIRES(c, TensorShapeUtils::IsVectorOrHigher(input_tensor.shape()), |
| 923 | errors::InvalidArgument( |
| 924 | "Tensor must be at least a vector, but saw shape: ", |
| 925 | input_tensor.shape().DebugString())); |
| 926 | TensorShape output_shape(input_tensor.shape()); |
| 927 | output_shape.RemoveDim(0); |
| 928 | OP_REQUIRES(c, element_shape.IsCompatibleWith(output_shape), |
| 929 | errors::InvalidArgument( |
| 930 | "Specified a list with shape ", element_shape.DebugString(), |
| 931 | " from a tensor with shape ", output_shape.DebugString())); |
| 932 | output_list.element_shape = element_shape; |
| 933 | |
| 934 | OP_REQUIRES(c, indices.NumElements() == input_tensor.shape().dim_size(0), |
| 935 | errors::InvalidArgument( |
| 936 | "Invalid number of rows in input tensor. Expected: ", |
| 937 | indices.NumElements(), |
| 938 | " Actual: ", input_tensor.shape().dim_size(0))); |
| 939 | |
| 940 | // Validate indices and resize output_list.tensors to fit the highest index. |
| 941 | { |
| 942 | int highest_index = -1; |
| 943 | for (int index = 0; index < indices.NumElements(); ++index) { |
| 944 | const int i = indices.flat<int32>()(index); |
| 945 | OP_REQUIRES( |
| 946 | c, i >= 0, |
| 947 | errors::InvalidArgument( |
| 948 | "Indices in TensorListScatter must all be non-negative.")); |
| 949 | OP_REQUIRES(c, num_elements == -1 || i < num_elements, |
| 950 | errors::InvalidArgument( |
| 951 | "TensorListScatter: Trying to scatter at index ", i, |
| 952 | " in list with size ", num_elements)); |
| 953 | if (i > highest_index) { |
| 954 | highest_index = i; |
| 955 | } |
| 956 | } |
| 957 | output_list.tensors().resize(std::max(highest_index + 1, num_elements), |
| 958 | Tensor(DT_INVALID)); |
| 959 | } |
| 960 | |
| 961 | OP_REQUIRES_OK(c, |
nothing calls this directly
no test coverage detected