class_docstr: Custom FX tracer for torchrec
  
  See `Torch.FX documentation <https://pytorch.org/docs/stable/fx.html>`_
  
  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
code: def trace(
      self,
      # pyre-ignore[2]: Missing parameter annotation [2]: Parameter `root` must have a type that does not contain `Any`
      root: Union[torch.nn.Module, Callable[..., Any]],
      concrete_args: Optional[Dict[str, Any]] = None,
  ) -> Graph:
      global _is_fx_tracing_flag
      old_is_fx_tracing_flag = _is_fx_tracing_flag
      _is_fx_tracing_flag = True
  
      try:
          # TODO(ivankobzarev): support DMP not only on the root level
          from torchrec.distributed.model_parallel import (
              DistributedModelParallel,
          )
  
          if isinstance(root, DistributedModelParallel):
              # In the case where the module is wrapped in DMP, you need to replace DMP's forward
              # call with a new signature, one with explicit args, because fx can't handle variable args.
              # Furthermore, we need to provide the `fn_root` argument because when tracing a function,
              # fx uses an empty module as the root (unless one is explicitly provided), which leads to
              # issues with path_of_module and named_buffers.
  
              # TODO(shababayub): This can be removed if we either stop supporting dmp wrapping
              # for fx trace or strip dmp name in named_modules path (much like named_buffers).
              if isinstance(root, torch.nn.Module):
                  for prefix, module in root.named_modules():
                      # TODO(T140754678): Remove this workaround to _fx_path
                      module._fx_path = prefix
  
              dmp = root
              graph = super().trace(
                  root=dmp_fx_trace_forward(dmp, self),
                  concrete_args=concrete_args,
              )
              self.root._dmp_wrapped_module = dmp._dmp_wrapped_module
          else:
              # Unwrapped dmp modules and composibility api will enter here.
              graph = super().trace(
                  root=root,
                  concrete_args=concrete_args,
              )
      finally:
          _is_fx_tracing_flag = old_is_fx_tracing_flag
      return graph
  
file_path: torchrec/torchrec/fx/tracer.py
imports: from typing import Any, Callable, Dict, List, Optional, Union
  import torch
  from torch.fx._compatibility import compatibility
  from torch.fx.graph import Graph
  from torch.fx.node import Argument
  from torchrec.distributed.types import LazyAwaitable, NoWait
  from torchrec.fx.utils import dmp_fx_trace_forward
label_desc: Use image node_id 3 for calling the Tracer obj's underlying member method code with example usage: obj.trace(root, concrete_args) and returns: graph