(
f: Callable[..., Value],
args: Tuple[LeafValue, ...],
guards: Set[Guard],
in_spec: Optional[TreeSpec] = None,
enable_functionalization: bool = True,
)
| 571 | |
| 572 | |
| 573 | def flattened_dispatch_trace( |
| 574 | f: Callable[..., Value], |
| 575 | args: Tuple[LeafValue, ...], |
| 576 | guards: Set[Guard], |
| 577 | in_spec: Optional[TreeSpec] = None, |
| 578 | enable_functionalization: bool = True, |
| 579 | ) -> Tuple[torch.fx.GraphModule, Value]: |
| 580 | if not isinstance(args, tuple): |
| 581 | raise TypeError(f"Expecting 'args' to be a tuple, got: {type(args)}") |
| 582 | |
| 583 | tracer = DispatchTracer() |
| 584 | |
| 585 | if enable_functionalization: |
| 586 | f = functionalize(f, remove="mutations_and_views") |
| 587 | tree_out = tracer.trace(f, concrete_args=args, in_spec=in_spec) |
| 588 | |
| 589 | name = type(f).__name__ if isinstance(f, torch.nn.Module) else f.__name__ |
| 590 | gm = torch.fx.GraphModule(tracer.root, tracer.graph, name) |
| 591 | |
| 592 | return (gm, tree_out) |
| 593 | |
| 594 | |
| 595 | @dataclass |
no test coverage detected