Shortcuts

torchrec.fx

Torchrec Tracer

Custom FX tracer for torchrec

See Torch.FX documentation

torchrec.fx.tracer

class torchrec.fx.tracer.Tracer(leaf_modules: Optional[List[str]] = None)

Bases: Tracer

Custom FX tracer for torchrec

See Torch.FX documentation

We create a custom FX tracer to trace torchrec based models. The custom tracer handles python generic types (i.e. NoWait[T], Awaitable[T]) and lower it to TorchScript if needed

create_arg(a: Any) Optional[Union[Tuple[Any, ...], List[Any], Dict[str, Any], slice, range, Node, str, int, float, bool, complex, dtype, Tensor, device, memory_format, layout, OpOverload]]

A method to specify the behavior of tracing when preparing values to be used as arguments to nodes in the Graph.

Adds support for the NoWait type in addition to the default tracer

Parameters:

a (Any) – The value to be emitted as an Argument in the Graph.

Returns:

The value a converted into the appropriate Argument

Return type:

Argument

graph: Graph
is_leaf_module(m: Module, module_qualified_name: str) bool

Override FX definition to include quantized embedding bags

module_stack: OrderedDict[str, Tuple[str, Any]]
node_name_to_scope: Dict[str, Tuple[str, type]]
path_of_module(mod: Module) str

Allows trace-ability of non registered modules. This is typically used for Table Batched Embeddings made to look like nn.EmbeddingBags

scope: Scope
trace(root: Union[Module, Callable[[...], Any]], concrete_args: Optional[Dict[str, Any]] = None) Graph

Note

Backwards-compatibility for this API is guaranteed.

torchrec.fx.tracer.is_fx_tracing() bool
torchrec.fx.tracer.symbolic_trace(root: Union[Module, Callable], concrete_args: Optional[Dict[str, Any]] = None, leaf_modules: Optional[List[str]] = None) GraphModule

Symbolic tracing API

Given an nn.Module or function instance root, this function will return a GraphModule constructed by recording operations seen while tracing through root.

concrete_args allows you to partially specialize your function, whether it’s to remove control flow or data structures.

Parameters:
  • root (Union[torch.nn.Module, Callable]) – Module or function to be traced and converted into a Graph representation.

  • concrete_args (Optional[Dict[str, any]]) – Inputs to be partially specialized

Returns:

a Module created from the recorded operations from root.

Return type:

GraphModule

Module contents

Torchrec Tracer

Custom FX tracer for torchrec

See Torch.FX documentation

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources