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

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

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