code: def create_fw_bw_graph(f, num_mapped_args, *args):
      mapped_xs = args[:num_mapped_args]
      pos_args = args[num_mapped_args:]
  
      # Note: We create "clean" environments for make_fx by suspending all dispatch keys
      # between Autograd and Python key. Currently, we only suspend functionalization but more can be
      # added when required. Will encounter two problems if we don't suspend functionalization:
      #
      # 1. make_fx fails to capture operations on input: the inputs are wrapped as _to_functional_tensor_wrapper,
      # but they will be unwrapped before entering ProxyTorchDispatchMode as part of the dispatching.
      # However, it's the outside wrapper that tracer creates proxies for. This casuses tracer fail to
      # fetch the proxy for the inputs and fail to capture any operations on them.
      #
      # 2. make_fx fails to capture output: the outputs after ProxyTorchDispatchMode are further
      # wrapped as FunctionalTensorWrapper in Functionalize key after return. However, the tracer
      # only associates the inner tensor with proxy in ProxyTorchDispatchMode. Therefore,
      # when creating the output node, it fails to associate the wrapped tensor with its proxy.
      # Instead, it will create _tensor_constant as output.
  
      with suspend_functionalization(), unset_functional_temporarily():
          with disable_proxy_modes_tracing():
  
              def _from_fun(t):
                  if isinstance(t, torch.Tensor):
                      if t.dtype != torch.bool:
                          return torch.empty_strided(
                              t.size(),
                              t.stride(),
                              dtype=t.dtype,
                              requires_grad=t.requires_grad,
                          )
                      else:
                          # clone of a functional tensor produces a functional tensor
                          # but we want to avoid it so we clone a non-functional version
                          maybe_unfunc_t = t
                          if isinstance(t, FunctionalTensor):
                              torch._sync(t)
                              maybe_unfunc_t = from_fun(t)
                          elif torch._is_functional_tensor(t):
                              # need to handle both types of functionalization here:
                              # these are the tensors that came from the user,
                              # which could be either FunctionalTensorWrapper or FunctionalTensor
                              torch._sync(t)
                              maybe_unfunc_t = (
                                  torch._from_functional_tensor(t)
                              )
                          return maybe_unfunc_t.clone()
                  return t
  
              example_xs = [
                  _from_fun(xs) for xs in _unstack_pytree(mapped_xs)[0]
              ]
  
              example_pos_args = [
                  _from_fun(arg)
                  if isinstance(arg, torch.Tensor)
                  else arg
                  for arg in pos_args
              ]
              example_flat_out = pytree.tree_map(
                  _from_fun, f(*example_xs, *example_pos_args)
              )
              if any(
                  not isinstance(out, torch.Tensor)
                  for out in example_flat_out
                  if out is not None
              ):
                  raise RuntimeError(
                      "Expect outputs of map only contains tensors or None. "
                      f"Got types {[type(out) for out in example_flat_out]}."
                  )
              example_grad = [
                  _from_fun(out) for out in example_flat_out
              ]
  
              fw_graph = make_fx(f)(*example_xs, *example_pos_args)
  
          def joint_f(*example_args):
              joint_mapped_args = example_args[:joint_num_mapped]
              args = example_args[joint_num_mapped:]
  
              mapped_input = joint_mapped_args[:num_mapped_args]
              mapped_grads = joint_mapped_args[num_mapped_args:]
  
              def fw_with_masks(*args):
                  fw_out = f(*args)
                  return fw_out, [
                      True
                      if isinstance(ret, torch.Tensor)
                      and ret.requires_grad
                      else False
                      for ret in fw_out
                  ]
  
              joint = create_joint(
                  fw_with_masks, aot_config=dummy_aot_config
              )
              _, grads = joint(
                  list(mapped_input) + list(args),
                  [
                      grad
                      for grad in mapped_grads
                      if grad is not None and grad.requires_grad
                  ],
              )
  
              # In order to keep map functional for backward graph,
              # we clone outputs that are aliasing inputs
              input_storage = {
                  StorageWeakRef(arg._typed_storage())
                  for arg in example_args
                  if isinstance(arg, torch.Tensor)
              }
  
              def maybe_clone(t):
                  if (
                      isinstance(t, torch.Tensor)
                      and StorageWeakRef(t._typed_storage())
                      in input_storage
                  ):
                      return t.clone()
                  return t
  
              return pytree.tree_map(maybe_clone, grads)
  
          joint_num_mapped = len(example_grad) + len(example_xs)
          joint_graph = make_fx(joint_f)(
              *example_xs, *example_grad, *example_pos_args
          )
          return fw_graph, joint_graph
  
imports: from contextlib import nullcontext
  import torch
  import torch.utils._pytree
  from torch._C import DispatchKey
  from torch._dispatch.python import suspend_functionalization
  from torch._functorch.aot_autograd import AOTConfig, create_joint, from_fun
  from torch._higher_order_ops.cond import _has_potential_branch_input_alias, _has_potential_branch_input_mutation, UnsupportedAliasMutationException
  from torch._ops import HigherOrderOperator
  from torch._subclasses.fake_tensor import FakeTensorMode
  from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode, unset_functional_temporarily
  from torch.fx.experimental.proxy_tensor import disable_proxy_modes_tracing, make_fx, ProxyTorchDispatchMode, track_tensor_tree
  from torch.multiprocessing.reductions import StorageWeakRef
label_desc: Use image node_id 1 for calling a global function with example usage: create_fw_bw_graph(f, num_mapped_args) and returns: fw_graph,  joint_graph, pytree, t, fw_out,  unknown, t, t, torch, maybe_unfunc_t