Executes a given callable `f` with a given tuple of arguments. During execution, Tensor operations are recorded in a fx.GraphModule, which is then returned. Args: f: A `nn.Module` or a Python function that implements an ML program. args: A tuple of arguments of any
(
f: Callable[..., Value],
args: Tuple[Value, ...],
)
| 705 | |
| 706 | |
| 707 | def dispatch_trace( |
| 708 | f: Callable[..., Value], |
| 709 | args: Tuple[Value, ...], |
| 710 | ) -> torch.fx.GraphModule: |
| 711 | """ |
| 712 | Executes a given callable `f` with a given tuple of arguments. During |
| 713 | execution, Tensor operations are recorded in a fx.GraphModule, which is then |
| 714 | returned. |
| 715 | |
| 716 | Args: |
| 717 | f: A `nn.Module` or a Python function that implements an ML program. |
| 718 | args: A tuple of arguments of any type to be used as inputs for the tracing run. |
| 719 | |
| 720 | Returns: |
| 721 | EXIR contained in a fx.GraphModule |
| 722 | """ |
| 723 | trace_func = f |
| 724 | guards = set() |
| 725 | if TORCHDYNAMO_ENABLED: |
| 726 | # Copying args is safer in case downstream implementations of trace_func mutate them |
| 727 | trace_func, guards = dynamo_trace(trace_func, args, False) |
| 728 | |
| 729 | # Copying args is safer in case downstream implementations of trace_func mutate them |
| 730 | trace_args, in_spec = pytree.tree_flatten(args) |
| 731 | |
| 732 | in_args = copy.deepcopy(tuple(trace_args)) |
| 733 | gm, tree_out = flattened_dispatch_trace( |
| 734 | trace_func, |
| 735 | in_args, |
| 736 | guards, |
| 737 | in_spec, |
| 738 | enable_functionalization=False, |
| 739 | ) |
| 740 | |
| 741 | _, out_spec = pytree.tree_flatten(tree_out) |
| 742 | |
| 743 | # pyre-fixme[16]: `GraphModule` has no attribute `in_spec`. |
| 744 | gm.in_spec = in_spec |
| 745 | # pyre-fixme[16]: `GraphModule` has no attribute `out_spec`. |
| 746 | gm.out_spec = out_spec |
| 747 | |
| 748 | # TODO (tmanlaibaatar) This is bit clowny, but our |
| 749 | # dispatch_trace sometimes creates unused node that |
| 750 | # breaks functionalization. it seems too much trouble |
| 751 | # to fix it properly since dispatch_trace will be deprecated soon. |
| 752 | # Basically dispatch_trace struggles on: |
| 753 | # def f(x: torch.Tensor) -> torch.Tensor: |
| 754 | # return torch.ones(6, dtype=x.dtype) |
| 755 | changed = gm.graph.eliminate_dead_code() |
| 756 | if changed: |
| 757 | gm.recompile() |
| 758 | |
| 759 | in_args = copy.deepcopy(tuple(trace_args)) |
| 760 | assert callable(gm) |
| 761 | |
| 762 | # This wrapper is used for preserving the stacktrace |
| 763 | # during second round of tracing. |
| 764 | # pyre-ignore |
no test coverage detected