Skip to content

trace_module

torch_tvarant.compiler.trace_module(module) → GraphModule

FX-trace module, inlining Linear / common activations.

Parameters:

Name Type Description Default
module Module

torch.nn.Module to trace.

 required

Returns:

Type Description
GraphModule

torch.fx.GraphModule: Graph with inlined leaf modules so the fuser
GraphModule

sees ATen ops (linear, relu, …).
Source code in torch_tvarant/compiler.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
def trace_module(module: nn.Module) -> torch.fx.GraphModule:
    """FX-trace ``module``, inlining Linear / common activations.

    Args:
        module: ``torch.nn.Module`` to trace.

    Returns:
        torch.fx.GraphModule: Graph with inlined leaf modules so the fuser
        sees ATen ops (``linear``, ``relu``, …).
    """
    tracer = TvarantTracer()
    graph = tracer.trace(module)
    return torch.fx.GraphModule(module, graph)