Wrapper around `body` that handles infeed queues and control deps.
(*inputs)
| 99 | return condition(*inputs) |
| 100 | |
| 101 | def body_wrapper(*inputs): |
| 102 | """Wrapper around `body` that handles infeed queues and control deps.""" |
| 103 | inputs = list(inputs) |
| 104 | |
| 105 | # Discards the dummy output added for arity-0 loops. |
| 106 | if input_arity == 0: |
| 107 | inputs = [] |
| 108 | |
| 109 | # Runs `body` with the dequeue_ops appended. |
| 110 | if infeed_queue: |
| 111 | number_of_shards = tpu_function.get_tpu_context().number_of_shards |
| 112 | if number_of_shards is None: |
| 113 | raise ValueError("Can't build training loop with infeed when there is " |
| 114 | "no tpu_shard_context. Are you building a loop or " |
| 115 | "graph directly rather than from inside tpu.rewrite, " |
| 116 | "tpu.batch_parallel, tpu.shard, or tpu.replicate?") |
| 117 | infeed_queue.set_number_of_shards(number_of_shards) |
| 118 | dequeue_ops = [d for d in infeed_queue.generate_dequeue_op()] |
| 119 | else: |
| 120 | dequeue_ops = [] |
| 121 | outputs = body(*(inputs + dequeue_ops)) |
| 122 | |
| 123 | # If the computation only returned one value, make it a tuple. |
| 124 | if not isinstance(outputs, (list, tuple)): |
| 125 | outputs = (outputs,) |
| 126 | |
| 127 | outputs = [ |
| 128 | o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o) |
| 129 | for o in outputs |
| 130 | ] |
| 131 | |
| 132 | # Separates the returned Operations and Tensors. |
| 133 | output_operations = [o for o in outputs if isinstance(o, ops.Operation)] |
| 134 | output_tensors = [o for o in outputs |
| 135 | if not isinstance(o, ops.Operation)] |
| 136 | |
| 137 | if outputs != output_tensors + output_operations: |
| 138 | raise ValueError( |
| 139 | "TPU training loop body must return zero or more Tensor values " |
| 140 | "followed by zero or more Operations.") |
| 141 | |
| 142 | output_types = [op.dtype for op in output_tensors] |
| 143 | if input_types != output_types: |
| 144 | raise TypeError( |
| 145 | "Mismatch between input types and output types for training loop " |
| 146 | "body: {} vs {}".format(input_types, output_types)) |
| 147 | |
| 148 | # Add the dequeue operations to output_operations to ensure they are run |
| 149 | # by the loop, even if the programmer's loop body does not use them. |
| 150 | output_operations += dequeue_ops |
| 151 | |
| 152 | # Add a dummy output, if needed. |
| 153 | if not output_tensors: |
| 154 | output_tensors = array_ops.constant(0) |
| 155 | |
| 156 | if output_operations: |
| 157 | # TODO(phawkins): in principle this is too restrictive since it serializes |
| 158 | # the training loop steps. In practice it does not matter since this loop |
nothing calls this directly
no test coverage detected