Shortcuts

torchrec.fx

Torchrec Tracer

Custom FX tracer for torchrec

See Torch.FX documentation

torchrec.fx.tracer

class torchrec.fx.tracer.Tracer

Bases: torch.fx._symbolic_trace.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, Node, str, int, float, bool, torch.dtype, torch.Tensor, torch.device, torch.memory_format, torch.layout]]

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: torch.fx.graph.Graph
torchrec.fx.tracer.symbolic_trace(root: Union[torch.nn.modules.module.Module, Callable], concrete_args: Optional[Dict[str, Any]] = None) torch.fx.graph_module.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