(
fn: Callable[..., Union[torch.Tensor, Tuple[torch.Tensor]]],
example_returns: Optional[List[torch.Tensor]] = None,
single_return: bool = False,
)
| 93 | |
| 94 | |
| 95 | def _make_submodule( |
| 96 | fn: Callable[..., Union[torch.Tensor, Tuple[torch.Tensor]]], |
| 97 | example_returns: Optional[List[torch.Tensor]] = None, |
| 98 | single_return: bool = False, |
| 99 | ) -> torch.fx.GraphModule: |
| 100 | if not hasattr(fn, "__tracing_inputs__"): |
| 101 | raise ExportError( |
| 102 | ExportErrorType.MISSING_PROPERTY, |
| 103 | f"Expect function '{fn.__name__}' to be decorated with tracing_context.", |
| 104 | ) |
| 105 | # pyre-ignore |
| 106 | args = fn.__tracing_inputs__ |
| 107 | # TODO(yidi): we don't want to enable here because we are not gonna use this code path in the future anyways |
| 108 | gm, _ = flattened_dispatch_trace(fn, args, set(), enable_functionalization=False) |
| 109 | output = next(iter(reversed(gm.graph.nodes))) |
| 110 | if example_returns: |
| 111 | internal_assert( |
| 112 | len(example_returns) == len(output.args[0]), |
| 113 | f"Eager mode of this {gm} returns {len(example_returns)} elements, but this graph returns {len(output.args[0])} elements", |
| 114 | ) |
| 115 | |
| 116 | if single_return: |
| 117 | # Force number of returned value to be 1. |
| 118 | internal_assert( |
| 119 | len(output.args[0]) == 1, |
| 120 | f"Graph {gm} should return just one element, but got {len(output.args[0])}", |
| 121 | ) |
| 122 | output.args = tuple(output.args[0]) |
| 123 | gm.recompile() |
| 124 | # pyre-fixme[16]: `GraphModule` has no attribute `__tracing_inputs__`. |
| 125 | gm.__tracing_inputs__ = args |
| 126 | return gm |
| 127 | |
| 128 | |
| 129 | def while_loop( |
no test coverage detected