code: def concrete_trace(
      root: Union[torch.nn.Module, Callable[..., Any]],
      concrete_args: Union[Dict[str, Any], Tuple],
      *,
      use_operator_patch: bool = True,
      operator_patch_backlist: List[str] | None = None,
      forward_function_name: str = "forward",
      check_args: Optional[Dict[str, Any]] = None,
      autowrap_leaf_function=None,
      autowrap_leaf_class=None,
      leaf_module: Tuple | None = None,
      fake_middle_class=None,
      dce=True,
      cpu_offload=False,
      trace_twice=False,
  ) -> GraphModule:
      """
      Concrete tracing API
  
      Given an ``nn.Module`` or function instance ``root`` and a dummy input `concrete_args`, this function will return a ``GraphModule``
      constructed by recording operations seen while tracing through ``root``.
  
      It has solved many problems compared to fx.symbolic_trace, and can execute on many third-party models.
  
      For example::
  
          def f(a, b):
              return a + b
  
          traced_f = concrete_trace(f, concrete_args={'a': 1, 'b': 2})
          # or `traced_f = concrete_trace(f, (1, 2))`
          assert traced_f(3, 4) == 7
  
          def f(x):
              out1, out2 = 0, 0
              for k, v in x.items():
                  out1 += k
                  out2 += v
              return out1, out2
          traced_f = concrete_trace(f, ({1: 1, 2: 2}, ))
          assert traced_f({2: 3, 4: 5}) == (6, 8)
  
      Note that we can only record static structure, so all the branches such as if-else or loop will be flattened::
  
          def f(x):
              out1, out2 = 0, 0
              for k, v in x.items():
                  out1 += k
                  out2 += v
              return out1, out2
          traced_f = concrete_trace(f, ({1: 1, 2: 2}, ))
          assert traced_f({2: 3, 4: 5, 6:7}) == (6, 8) # not (12, 15)
  
          # traced code like:
          def traced_f(self, x):
              out1, out2 = 0, 0
              items = x.items()
  
              # for loop
              iter = iter(items)
  
              # first loop content
              items0 = next(iter)
              out1 += items0[0]
              out2 += items0[1]
  
              # second loop content
              items1 = next(iter)
              out1 += items1[0]
              out2 += items1[1]
  
              return (out1, out2)
  
      If you want to trace 'is', 'is not', 'in' or 'not in' in your module, you can set use_function_patch to True::
  
          def f(x, y):
              if x is None:
                  return y
              else:
                  return x - y
          # traced_f = concrete_trace(f, (None, 1)) # bad
          traced_f = concrete_trace(f, (None, 1), use_function_patch=True) # f should exist in a file.
  
      If you have a function/method that should be treated as a leaf function but not trace into it, use autowrap_leaf_function to mark it::
  
          def leaf_op(x, y, z):
              # if not treated as a leaf function, then only 1 branch will exist.
              if x > 0:
                  return y + z
              else:
                  return y - z
  
          def f(x):
              return leaf_op(x, 3, 2)
  
          traced_f = concrete_trace(f, (1, ), autowrap_leaf_function = {
              leaf_op: ([], False, None), **ConcreteTracer.default_autowrap_leaf_function})
          assert traced_f(1) == 5 and traced_f(-1) == 1
  
      If you have a class that should be treated as a leaf class, use autowrap_leaf_class to mark it::
  
          class leaf_clz:
              def __init__(self, a, b):
                  self.c = a + b
  
          def f(x, y):
              return leaf_clz(x, y)
  
          traced_f = concrete_trace(f, (1, 2), autowrap_leaf_class = {
              leaf_clz: ([], False), **ConcreteTracer.default_autowrap_leaf_class})
          assert isinstance(traced_f(3, 4), leaf_clz) and traced_f(3, 4).c == 7
  
      Args:
          root (Union[torch.nn.Module, Callable]): Module or function to be traced and converted into a Graph representation.
          concrete_args (Union[Dict[str, Any], Tuple]): Dummy inputs to do concrete trace.
  
          use_function_patch (bool): Use operator patcher recursively on function calls. Operator patcher will re-compile the function and
              translate '{} is {}' into 'operator.is_({}, {})', then we can treat 'is', 'is not', 'in' and 'not in' as function calls.
  
          operator_patch_backlist (List[str]): Blacklist of the operator patcher.
  
          autowrap_leaf_function (Dict[Any, Tuple[List[Tuple[Union[ModuleType, Type], str]], bool, Optional[Callable]]]): Leaf function dict,
              such as 'add' or 'torch.xxx'. You can add your own leaf functions.
  
              The struct of dict is: leaf_function: ([(module_path, module_name)], force_to_trace, replace_to_function).
                  (module_path, module_name): The place the function exists. Such as torch.meshgrid, there are `torch.meshgrid`,
                      'torch.functional.meshgrid', 'torch._C._VariableFunctions.meshgrid', we should wrap them all.
                  force_to_trace: If set to false, the function will only be traced if input relates to concrete_args.
                      Such as 'torch.rand', we should trace it even if it doesn't relate to concrete_args.
                  replace_to_function: If not `None`, we will use it to replace the original function in traced code.
                      Such as ModuleList.__getitem__, we can use operator.getitem to replace it.
  
          default_autowrap_leaf_class (Dict[Type, Tuple[List[Tuple[Union[ModuleType, Type], str]], bool]]): Leaf class dict, such as 'int',
              'range' or 'zip'. You can add your own leaf functions such as 'torch.finfo' or 'modeling_outputs.SequenceClassifierOutput'.
  
              The struct of dict is: leaf_class: ([(module_path, module_name)], is_iterator_class).
                  is_iterator_class: Is the class init from an iterator. Only 'tuple', 'list', 'set' or 'dict' needs to set it to True.
  
          cpu_offload (bool): Whether to offload the module to CPU during tracing. If set to True, the traced code will be executed on GPU,
              but is offloaded to CPU afterward. This is useful for reducing memory usage during tracing, but may cause performance issues.
              If set to False, there will be no offloading during tracing, but the traced code will be executed on default device.
  
      Returns:
          fx.GraphModule: a Module created from the recorded operations from ``root``.
      """
      tracer = ConcreteTracer(cpu_offload=cpu_offload)
      is_training = root.training
      root.eval()
  
      graph = tracer.trace(
          root,
          autowrap_leaf_function=autowrap_leaf_function,
          autowrap_leaf_class=autowrap_leaf_class,
          leaf_module=leaf_module,
          fake_middle_class=fake_middle_class,
          concrete_args=concrete_args,
          use_operator_patch=use_operator_patch,
          operator_patch_backlist=operator_patch_backlist,
          forward_function_name=forward_function_name,
      )
  
      if trace_twice:
          graph_check = tracer.trace(
              root,
              autowrap_leaf_function=autowrap_leaf_function,
              autowrap_leaf_class=autowrap_leaf_class,
              leaf_module=leaf_module,
              fake_middle_class=fake_middle_class,
              concrete_args=concrete_args,
              use_operator_patch=use_operator_patch,
              operator_patch_backlist=operator_patch_backlist,
              forward_function_name=forward_function_name,
          )
          # compare to check equal
          assert len(graph.nodes) == len(
              graph_check.nodes
          ), f"number nodes: {len(graph.nodes)} vs {len(graph_check.nodes)}"
          for node_a, node_b in zip(graph.nodes, graph_check.nodes):
              node_a: Node
              node_b: Node
              target_a = node_a.target
              target_b = node_b.target
              if node_a.op == "get_attr" and node_a.name.startswith(
                  "_tensor_constant"
              ):
                  assert (
                      node_b.op == "get_attr"
                      and node_b.name.startswith("_tensor_constant")
                  )
                  assert torch.equal(
                      getattr(root, node_a.name),
                      getattr(root, node_b.name),
                  )
              elif (
                  node_a.op == "call_function"
                  and isinstance(target_a, Callable)
                  and target_a.__name__ == "apply"
                  and hasattr(target_a, "__self__")
                  and issubclass(
                      target_a.__self__, torch.autograd.Function
                  )
              ):
                  assert (
                      node_b.op == "call_function"
                      and isinstance(target_b, Callable)
                      and target_b.__name__ == "apply"
                      and hasattr(target_b, "__self__")
                      and issubclass(
                          target_b.__self__, torch.autograd.Function
                      )
                  )
              else:
                  assert (
                      node_a.op == node_b.op and target_a == target_b
                  ), f"op: {node_a.op} vs {node_b.op}, target: {target_a} vs {target_b}"
  
      with MagicMethodPatcher():
          name = (
              root.__class__.__name__
              if isinstance(root, torch.nn.Module)
              else root.__name__
          )
          traced = GraphModule(tracer.root, graph, name)
  
          if dce:
              with _Patcher() as patcher:
                  patcher.patch_method(
                      Node,
                      "is_impure",
                      node_is_impure_wrapper,
                      deduplicate=False,
                  )
                  traced.graph.eliminate_dead_code()
              traced.recompile()  # this need to be done in MagicMethodPatcher context
  
      # TODO: better infomation
      if check_args is not None:
          assert root(**check_args) == traced(**check_args)
  
      if is_training:
          root.train()
  
      return traced
  
imports: from __future__ import annotations
  import collections
  import copy
  import sys
  import inspect
  import logging
  import operator
  import functools
  import builtins
  from itertools import chain
  from types import BuiltinMethodType, FunctionType, MethodDescriptorType, MethodType, MethodWrapperType, ModuleType
  from typing import Any, Dict, Iterable, Iterator, Optional, Set, Tuple, Type, List, Callable, Union
  from contextlib import contextmanager
  import torch
  from torch._C import ScriptObject
  from torch.nn.modules.container import Sequential, ModuleList, ModuleDict, ParameterList, ParameterDict
  from torch.utils._pytree import tree_map
  import torch.fx
  from torch.fx import GraphModule
  from torch.fx._compatibility import compatibility
  from torch.fx._symbolic_trace import _Patcher, _proxyable_classes
  from torch.fx.graph import Graph
  from torch.fx.node import Target, Node, Argument, _side_effectful_functions
  from torch.fx.proxy import TracerBase
  from torch.fx.operator_schemas import check_for_mutable_operation
  from .None import concrete_proxy
  from .operator_patcher import OperatorPatcherContext
  from .utils import _orig_module_call, _orig_module_getattr, _orig_module_getattribute, _orig_agfunc_apply, _orig_torch_assert, _orig_type, _orig_isinstance, _orig_issubclass, _orig_getattr, _orig_range, _orig_int, _orig_bool, _orig_tuple, _orig_list, _orig_set, _orig_frozenset, _orig_dict, _orig_map, _orig_zip, _orig_enumerate, _orig_slice, _orig_reversed, _orig_torch_size, _orig_torch_finfo, _orig_len, _orig_not, _orig_is, _orig_is_not, _orig_contains, _orig_index, _orig_all, _orig_min, _orig_max, _orig_node_is_impure
label_desc: Use image node_id 13 for calling a global function with example usage: concrete_trace(root, concrete_args) and returns: traced