Utility class for representing a list of output tensors that are grouped as a single named output.
| 511 | // Utility class for representing a list of output tensors that are |
| 512 | // grouped as a single named output. |
| 513 | class OpOutputList { |
| 514 | public: |
| 515 | typedef OpArgIterator<OpOutputList, const Tensor*> Iterator; |
| 516 | OpOutputList() : ctx_(nullptr), start_(0), stop_(0) {} |
| 517 | OpOutputList(OpKernelContext* ctx, int start, int stop) |
| 518 | : ctx_(ctx), start_(start), stop_(stop) {} |
| 519 | OpOutputList& operator=(const OpOutputList& other) = default; |
| 520 | Tensor* operator[](int i); |
| 521 | bool required(int i) const; |
| 522 | DataType expected_output_dtype(int i) const; |
| 523 | Status allocate(int i, const TensorShape& shape, Tensor** output); |
| 524 | void set(int i, const Tensor& tensor); |
| 525 | void set_ref(int i, mutex* mu, Tensor* tensor_for_ref); |
| 526 | int size() const { return stop_ - start_; } |
| 527 | Iterator begin() const { return Iterator(this, 0); } |
| 528 | Iterator end() const { return Iterator(this, size()); } |
| 529 | |
| 530 | private: |
| 531 | OpKernelContext* ctx_; // not owned |
| 532 | int start_; |
| 533 | int stop_; |
| 534 | }; |
| 535 | |
| 536 | // Holds a tensor or tensor reference. For tensor references, we need |
| 537 | // a mutex to prevent concurrent access to the tensor. |