For a given node, return the dataclass representing its output values. [NOTE: Multiple outputs] We handle aggregates differently than FX. For FX, it looks like: x = call_function("multiple_return", ...) element0 = call_function(getitem, x, 0) foo
(self, node: torch.fx.Node)
| 1115 | ] |
| 1116 | |
| 1117 | def serialize_outputs(self, node: torch.fx.Node) -> List[Argument]: |
| 1118 | """For a given node, return the dataclass representing its output values. |
| 1119 | |
| 1120 | [NOTE: Multiple outputs] We handle aggregates differently than FX. For |
| 1121 | FX, it looks like: |
| 1122 | |
| 1123 | x = call_function("multiple_return", ...) |
| 1124 | element0 = call_function(getitem, x, 0) |
| 1125 | foo = call_function("use_output", element0) |
| 1126 | |
| 1127 | We do not want the intermediate `getitem` call, so our serialized thing looks like: |
| 1128 | |
| 1129 | element0, element1, element2 = call_function("multiple_return", ...) |
| 1130 | foo = call_function("use_output", element0) |
| 1131 | |
| 1132 | We want names to be consistent across these two schemes, so that we can |
| 1133 | mostly reuse the names coming from FX. This function computes a mapping from |
| 1134 | the FX representation to our representation, preserving the names. |
| 1135 | """ |
| 1136 | assert node.op == "call_function" and isinstance( |
| 1137 | node.target, torch._ops.OpOverload |
| 1138 | ) |
| 1139 | |
| 1140 | assert isinstance(node.target, torch._ops.OpOverload) |
| 1141 | returns = node.target._schema.returns |
| 1142 | |
| 1143 | if len(returns) == 0: |
| 1144 | return [] |
| 1145 | |
| 1146 | meta_val = node.meta["val"] |
| 1147 | |
| 1148 | # Check single value return |
| 1149 | if _is_single_tensor_list_return(node.target): |
| 1150 | # e.g "-> Tensor[]" |
| 1151 | tensor_args = [] |
| 1152 | for idx, meta in enumerate(meta_val): |
| 1153 | user_node = _output_node_at_index(node, idx) |
| 1154 | name = ( |
| 1155 | user_node.name |
| 1156 | if user_node is not None |
| 1157 | else f"{node.name}_unused_{idx}" |
| 1158 | ) |
| 1159 | tensor_args.append(self.serialize_tensor_output(name, meta)) |
| 1160 | return [Argument.create(as_tensors=tensor_args)] |
| 1161 | elif len(returns) == 1: |
| 1162 | return [self.serialize_output(node.name, meta_val)] |
| 1163 | |
| 1164 | # There are a two possibilities at this point: |
| 1165 | # - This operator returns a tuple of Tensors, e.g. "-> (Tensor, Tensor)" |
| 1166 | # - This operator returns a tuple of mixed of Tensor and Tensors, e.g. "-> (Tensor, Tensor[])" |
| 1167 | # |
| 1168 | # Either way, start by gathering a list of TensorArguments with the correct names. |
| 1169 | # For consistent naming with FX, consult the downstream `getitem` node and |
| 1170 | # make sure our outputs have the same name. |
| 1171 | |
| 1172 | output_arguments = [] |
| 1173 | for idx, (meta, return_schema) in enumerate(zip(meta_val, returns)): |
| 1174 | if meta is None: |
no test coverage detected