Torchrec Tracer

Custom FX tracer for torchrec

See Torch.FX documentation


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


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


The value a converted into the appropriate Argument

Return type:


is_leaf_module(m: Module, module_qualified_name: str) bool

Override FX definition to include quantized embedding bags

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

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


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.

  • 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


a Module created from the recorded operations from root.

Return type:


Module contents

Torchrec Tracer

Custom FX tracer for torchrec

See Torch.FX documentation


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources