Validates non-flat outputs, add backs device assignments and other attrs. Args: outputs: Output from `computation` inside `tpu.rewrite`. Returns: Tensors extracted from outputs and an empty list because Operations are not allowed in non-flat outputs..
(outputs)
| 1128 | |
| 1129 | |
| 1130 | def _postprocess_non_flat_outputs(outputs): |
| 1131 | """Validates non-flat outputs, add backs device assignments and other attrs. |
| 1132 | |
| 1133 | Args: |
| 1134 | outputs: Output from `computation` inside `tpu.rewrite`. |
| 1135 | |
| 1136 | Returns: |
| 1137 | Tensors extracted from outputs and an empty list because Operations are not |
| 1138 | allowed in non-flat outputs.. |
| 1139 | """ |
| 1140 | |
| 1141 | # Flatten output items. |
| 1142 | flat_outputs = nest.flatten(outputs) |
| 1143 | |
| 1144 | # Convert all non-Operation outputs to Tensors. |
| 1145 | for i, o in enumerate(flat_outputs): |
| 1146 | if isinstance(o, ops.Operation): |
| 1147 | raise ValueError( |
| 1148 | "tpu.rewrite does not support Operation as return value in non-flat " |
| 1149 | "output structure. You can set returned Operations as control " |
| 1150 | "dependencies of returned Tensors so Operations are triggered when " |
| 1151 | 'Tensors are evaluated. Operation found: "%s"' % o.name) |
| 1152 | |
| 1153 | try: |
| 1154 | o = ops.convert_to_tensor(o) |
| 1155 | except Exception as e: |
| 1156 | raise ValueError( |
| 1157 | "TPU function return values must all either be Operations or " |
| 1158 | 'convertible to Tensors. Got error: "%s"' % str(e)) |
| 1159 | |
| 1160 | # Wraps outputs in Identity ops. Otherwise a replicated input copied |
| 1161 | # straight to an output would bypass the replicate(). This would be bad |
| 1162 | # because the TPUReplicatedInput/TPUReplicatedOutput operator would not |
| 1163 | # be rewritten away, leading to a runtime error. |
| 1164 | # TODO(phawkins): extend the rewrite to elide these nodes instead. |
| 1165 | with ops.device(core(0)): |
| 1166 | o = array_ops.identity(o) |
| 1167 | # pylint: disable=protected-access |
| 1168 | o.op._set_attr("_tpu_output_identity", attr_value_pb2.AttrValue(b=True)) |
| 1169 | # pylint: enable=protected-access |
| 1170 | flat_outputs[i] = array_ops.identity(o) |
| 1171 | |
| 1172 | # All flat_outputs are Tensors, and no Operations. |
| 1173 | return flat_outputs, [] |
| 1174 | |
| 1175 | |
| 1176 | def split_compile_and_shard(computation, |