code: def create_functionalized_rng_ops_wrapper(
      func, args, trace_joint=True
  ) -> Any:
      # Functionalization of rng ops changes the calling convention of the joint graph.
      # It goes from (primals, tangents) to (seed, offset, primals, tangents)
      # At runtime, we pass on the current seed and offset. This is hidden from
      # the user.
      fake_mode = detect_fake_mode()
      if fake_mode is None:
          fake_mode = nullcontext()
  
      def override_get_rng_state(
          device: Union[int, str, torch.device] = "cuda"
      ):
          out = PhiloxStateTracker.get_state_as_tensor()
          return out
  
      def override_set_rng_state(
          x, device: Union[int, str, torch.device] = "cuda"
      ):
          PhiloxStateTracker.set_state_from_tensor(x)
  
      def append_rng_offsets(args):
          if trace_joint:
              # args signature before: Tuple(fwd_outputs), Tuple(bwd_outputs)
              # args signature after: Tuple(fwd_outputs, new_fwd_rng_offset), Tuple(bwd_offset, new_bwd_rng_offset)
              return (
                  (
                      *args[0],
                      PhiloxStateTracker.get_updated_fwd_offset(),
                  ),
                  (
                      *args[1],
                      PhiloxStateTracker.get_updated_bwd_offset(),
                  ),
              )
          else:
              # args signature before: Tuple(fwd_outputs)
              # args signature after: Tuple(fwd_outputs, new_fwd_rng_offset)
              return (
                  *args,
                  PhiloxStateTracker.get_updated_fwd_offset(),
              )
  
      def traced_joint(
          primals,
          tangents,
          fwd_seed,
          fwd_base_offset,
          bwd_seed,
          bwd_base_offset,
      ):
          with patch(
              "torch.cuda.get_rng_state", override_get_rng_state
          ), patch("torch.cuda.set_rng_state", override_set_rng_state):
              return append_rng_offsets(func(primals, tangents))
  
      def traced_forward(*primals_fwd_seed_fwd_base_offset):
          # The signature is (*primals, seed, offset)
          with patch(
              "torch.cuda.get_rng_state", override_get_rng_state
          ), patch("torch.cuda.set_rng_state", override_set_rng_state):
              return append_rng_offsets(
                  func(*primals_fwd_seed_fwd_base_offset[:-2])
              )
  
      if trace_joint:
          # Get the current seed and offset to setup tracing.
          (
              fwd_seed,
              fwd_base_offset,
          ) = CUDARngStateHelper.get_torch_state_as_tuple(fake_mode)
          (
              bwd_seed,
              bwd_base_offset,
          ) = CUDARngStateHelper.get_torch_state_as_tuple(fake_mode)
          PhiloxStateTracker.record_state(
              fwd_seed, fwd_base_offset, "forward"
          )
          PhiloxStateTracker.record_state(
              bwd_seed, bwd_base_offset, "backward"
          )
          return traced_joint, (
              *args,
              fwd_seed,
              fwd_base_offset,
              bwd_seed,
              bwd_base_offset,
          )
      else:
          # Get the current seed and offset to setup tracing.
          (
              fwd_seed,
              fwd_base_offset,
          ) = CUDARngStateHelper.get_torch_state_as_tuple(fake_mode)
          PhiloxStateTracker.record_state(
              fwd_seed, fwd_base_offset, "forward"
          )
          return traced_forward, (*args, fwd_seed, fwd_base_offset)
  
imports: import warnings
  from contextlib import nullcontext
  from functools import wraps
  from typing import Any, Callable, List, Tuple, Union
  from unittest.mock import patch
  import torch
  import torch.fx.traceback
  import torch.utils._pytree
  from torch import Tensor
  from torch._decomp.decompositions_for_rng import PhiloxStateTracker
  from torch._guards import detect_fake_mode
  from torch._prims_common import CUDARngStateHelper
  from torch._subclasses.functional_tensor import FunctionalTensorMode
  from torch.fx import Interpreter
  from torch.fx.experimental.symbolic_shapes import definitely_false, sym_eq
  from torch.nn.utils import stateless
  from ..None import config
  from .collect_metadata_analysis import run_functionalized_fw_and_collect_metadata
  from .functional_utils import from_fun, has_data_mutation, has_metadata_mutation, is_fun, sync_functional_tensor, to_fun
  from .logging_utils import setup_stacktrace_preservation_hooks
  from .schemas import AOTConfig, MutationType, OutputType, SubclassMeta, SubclassTracingInfo, ViewAndMutationMeta
  from .subclass_utils import create_subclass_meta, requires_subclass_dispatch, unwrap_tensor_subclasses, wrap_tensor_subclasses_maybe_joint
  from .utils import maybe_to_fresh_input
label_desc: Use image node_id 4 for calling a global function with example usage: create_functionalized_rng_ops_wrapper(func, args, trace_joint) and returns: out, traced_joint,  unknown,  fwd_seed,  fwd_base_offset,  bwd_seed,  bwd_base_offset, traced_forward,  unknown,  fwd_seed,  fwd_base_offset, unknown,  PhiloxStateTracker,  unknown,  PhiloxStateTracker, unknown,  PhiloxStateTracker, append_rng_offsets, append_rng_offsets