Generate a host-side Op to enqueue a tuple to the queue. If device is None the inputs are all required to have the same device specification, and the enqueue Op is colocated with inputs[0]. Otherwise the enqueue Op is placed on 'device'. Args: inputs: a list of Tensors with t
(self,
inputs,
name_prefix,
index,
device=None,
tpu_ordinal=-1)
| 497 | dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name) |
| 498 | |
| 499 | def _generate_enqueue_op(self, |
| 500 | inputs, |
| 501 | name_prefix, |
| 502 | index, |
| 503 | device=None, |
| 504 | tpu_ordinal=-1): |
| 505 | """Generate a host-side Op to enqueue a tuple to the queue. |
| 506 | |
| 507 | If device is None the inputs are all required to have the same |
| 508 | device specification, and the enqueue Op is colocated with |
| 509 | inputs[0]. Otherwise the enqueue Op is placed on 'device'. |
| 510 | |
| 511 | Args: |
| 512 | inputs: a list of Tensors with the types and shapes of the tuple elements. |
| 513 | name_prefix: the base name for the Op. |
| 514 | index: the shard index, used to uniquify the Op name. |
| 515 | device: device to place the Op on, or None if it should be |
| 516 | colocated with the inputs. |
| 517 | tpu_ordinal: ordinal of the TPU device on the host to use for |
| 518 | infeed if device is a CPU device. Should be set to -1 if device |
| 519 | is a TPU device. |
| 520 | |
| 521 | Returns: |
| 522 | An Op corresponding to a shard of infeed enqueued at the host, |
| 523 | suitable for use within a replicated block. |
| 524 | |
| 525 | Raises: |
| 526 | ValueError: if device is None and inputs do not all have the |
| 527 | same device specification. |
| 528 | """ |
| 529 | full_name = "%s/%d" % (name_prefix, index) |
| 530 | shapes = [t.shape for t in inputs] |
| 531 | if device is None: |
| 532 | devices = [t.device for t in inputs] |
| 533 | for i in xrange(1, self.number_of_tuple_elements): |
| 534 | if devices[0] != devices[i]: |
| 535 | raise ValueError( |
| 536 | "input devices for shard %d are %s, but should all be the same" % |
| 537 | (index, str(devices))) |
| 538 | with ops.colocate_with(inputs[0]): |
| 539 | return tpu_ops.infeed_enqueue_tuple( |
| 540 | inputs=inputs, |
| 541 | shapes=shapes, |
| 542 | name=full_name, |
| 543 | device_ordinal=tpu_ordinal) |
| 544 | else: |
| 545 | with ops.device(device): |
| 546 | return tpu_ops.infeed_enqueue_tuple( |
| 547 | inputs=inputs, |
| 548 | shapes=shapes, |
| 549 | name=full_name, |
| 550 | device_ordinal=tpu_ordinal) |
| 551 | |
| 552 | def generate_enqueue_ops(self, |
| 553 | sharded_inputs, |
no test coverage detected