| 193 | } |
| 194 | |
| 195 | void FIFOQueue::TryDequeueMany(int num_elements, OpKernelContext* ctx, |
| 196 | bool allow_small_batch, |
| 197 | CallbackWithTuple callback) { |
| 198 | if (!specified_shapes()) { |
| 199 | ctx->SetStatus(errors::InvalidArgument( |
| 200 | "FIFOQueue's DequeueMany and DequeueUpTo require the " |
| 201 | "components to have specified shapes.")); |
| 202 | callback(Tuple()); |
| 203 | return; |
| 204 | } |
| 205 | if (num_elements == 0) { |
| 206 | Tuple tuple; |
| 207 | tuple.reserve(num_components()); |
| 208 | for (int i = 0; i < num_components(); ++i) { |
| 209 | // TODO(josh11b,misard): Switch to allocate_output(). Problem is |
| 210 | // this breaks the abstraction boundary since we don't *really* |
| 211 | // know if and how the Tensors in the tuple we pass to callback |
| 212 | // correspond to the outputs of *ctx. For example, the |
| 213 | // ReaderRead Op uses TryDequeue() to get a filename out of a |
| 214 | // queue that is used internally by the reader and is not |
| 215 | // associated with any output of the ReaderRead. |
| 216 | // mrry@ adds: |
| 217 | // Maybe we need to pass a std::function<Tensor*(...)> (or |
| 218 | // better signature) that calls the appropriate allocator |
| 219 | // function in addition to ctx? (Or support a shim Allocator |
| 220 | // that has an internal OpKernelContext*, and dispatches to the |
| 221 | // appropriate method?) |
| 222 | // misard@ adds: |
| 223 | // I don't see that a std::function would help. The problem is |
| 224 | // that at this point (allocation time) the system doesn't know |
| 225 | // what is going to happen to the element read out of the |
| 226 | // queue. As long as we keep the generality that TensorFlow Ops |
| 227 | // do their own dynamic allocation in arbitrary C++ code, we |
| 228 | // need to preserve robustness to allocating output Tensors with |
| 229 | // the 'wrong' attributes, and fixing up with a copy. The only |
| 230 | // improvement I can see here in the future would be to support |
| 231 | // an optimized case where the queue 'knows' what attributes to |
| 232 | // use, and plumbs them through here. |
| 233 | Tensor element; |
| 234 | Status status = ctx->allocate_temp(component_dtypes_[i], |
| 235 | ManyOutShape(i, 0), &element); |
| 236 | if (!status.ok()) { |
| 237 | ctx->SetStatus(status); |
| 238 | callback(Tuple()); |
| 239 | return; |
| 240 | } |
| 241 | tuple.emplace_back(element); |
| 242 | } |
| 243 | callback(tuple); |
| 244 | return; |
| 245 | } |
| 246 | |
| 247 | CancellationManager* cm = ctx->cancellation_manager(); |
| 248 | CancellationToken token = cm->get_cancellation_token(); |
| 249 | bool already_cancelled; |
| 250 | { |
| 251 | mutex_lock l(mu_); |
| 252 | already_cancelled = !cm->RegisterCallback( |
no test coverage detected