| 190 | } |
| 191 | |
| 192 | Status ForwardInputOrCreateNewList(OpKernelContext* c, int32 input_index, |
| 193 | int32 output_index, |
| 194 | const TensorList& input_list, |
| 195 | TensorList** output_list) { |
| 196 | // Attempt to forward the input tensor to the output if possible. |
| 197 | std::unique_ptr<Tensor> maybe_output = c->forward_input( |
| 198 | input_index, output_index, DT_VARIANT, TensorShape{}, |
| 199 | c->input_memory_type(input_index), AllocatorAttributes()); |
| 200 | Tensor* output_tensor; |
| 201 | if (maybe_output != nullptr && maybe_output->dtype() == DT_VARIANT && |
| 202 | maybe_output->NumElements() == 1) { |
| 203 | output_tensor = maybe_output.get(); |
| 204 | TensorList* tmp_out = output_tensor->scalar<Variant>()().get<TensorList>(); |
| 205 | if (tmp_out == nullptr) { |
| 206 | return errors::InvalidArgument( |
| 207 | "Expected input ", input_index, " to be a TensorList but saw ", |
| 208 | output_tensor->scalar<Variant>()().TypeName()); |
| 209 | } |
| 210 | if (tmp_out->RefCountIsOne()) { |
| 211 | // Woohoo, forwarding succeeded! |
| 212 | c->set_output(output_index, *output_tensor); |
| 213 | *output_list = tmp_out; |
| 214 | return Status::OK(); |
| 215 | } |
| 216 | } |
| 217 | |
| 218 | // If forwarding is not possible allocate a new output tensor and copy |
| 219 | // the `input_list` to it. |
| 220 | AllocatorAttributes attr; |
| 221 | attr.set_on_host(true); |
| 222 | TF_RETURN_IF_ERROR( |
| 223 | c->allocate_output(output_index, {}, &output_tensor, attr)); |
| 224 | output_tensor->scalar<Variant>()() = input_list.Copy(); |
| 225 | |
| 226 | *output_list = output_tensor->scalar<Variant>()().get<TensorList>(); |
| 227 | return Status::OK(); |
| 228 | } |
| 229 | |
| 230 | class EmptyTensorList : public OpKernel { |
| 231 | public: |
no test coverage detected