Validates non-flat outputs, add backs device assignments and other attrs. Args: outputs: Output from `computation` inside `tpu.rewrite`. Returns: Tensors and Operations extracted from outputs.
(outputs)
| 1067 | |
| 1068 | |
| 1069 | def _postprocess_flat_outputs(outputs): |
| 1070 | """Validates non-flat outputs, add backs device assignments and other attrs. |
| 1071 | |
| 1072 | Args: |
| 1073 | outputs: Output from `computation` inside `tpu.rewrite`. |
| 1074 | |
| 1075 | Returns: |
| 1076 | Tensors and Operations extracted from outputs. |
| 1077 | """ |
| 1078 | # Following code segment is to preserve legacy behavior. Previously we only |
| 1079 | # supported flat outputs and thus for consistency it was nice to convert even |
| 1080 | # single element into a tuple. But now that we support arbitrary output |
| 1081 | # structure, this is no longer necessary. |
| 1082 | # TODO(b/121383831): Migrate all legacy use cases and delete this special |
| 1083 | # case. |
| 1084 | # If the computation returns `None`, make it an empty tuple. |
| 1085 | if outputs is None: |
| 1086 | outputs = tuple() |
| 1087 | # If the computation only returned one value, makes it a tuple. |
| 1088 | if not isinstance(outputs, collections_abc.Sequence): |
| 1089 | outputs = (outputs,) |
| 1090 | |
| 1091 | # Append `no_op` here so that fetching any return value of this function |
| 1092 | # will trigger TPUExecute node. |
| 1093 | outputs += (control_flow_ops.no_op(),) |
| 1094 | try: |
| 1095 | with ops.device(core(0)): |
| 1096 | outputs = [ |
| 1097 | o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o) |
| 1098 | for o in outputs |
| 1099 | ] |
| 1100 | except Exception as e: |
| 1101 | raise ValueError( |
| 1102 | "TPU function return values must all either be Operations or " |
| 1103 | "convertible to Tensors. Got '%s'" % str(e)) |
| 1104 | |
| 1105 | # Separates the returned Operations and Tensors. |
| 1106 | output_operations = [o for o in outputs if isinstance(o, ops.Operation)] |
| 1107 | output_tensors = [o for o in outputs if not isinstance(o, ops.Operation)] |
| 1108 | |
| 1109 | if outputs != output_tensors + output_operations: |
| 1110 | raise ValueError( |
| 1111 | "TPU functions must return zero-or more Tensor values followed by " |
| 1112 | "zero or more Operations.") |
| 1113 | |
| 1114 | # Wraps outputs in Identity ops. Otherwise a replicated input copied |
| 1115 | # straight to an output would bypass the replicate(). This would be bad |
| 1116 | # because the TPUReplicatedInput/TPUReplicatedOutput operator would not |
| 1117 | # be rewritten away, leading to a runtime error. |
| 1118 | # TODO(phawkins): extend the rewrite to elide these nodes instead. |
| 1119 | new_output_tensors = [] |
| 1120 | for t in output_tensors: |
| 1121 | with ops.device(t.device if t.device else core(0)): |
| 1122 | o = array_ops.identity(t) |
| 1123 | # pylint: disable=protected-access |
| 1124 | o.op._set_attr("_tpu_output_identity", attr_value_pb2.AttrValue(b=True)) |
| 1125 | # pylint: enable=protected-access |
| 1126 | new_output_tensors.append(o) |