code: def aot_dispatch_subclass(
      flat_fn_maybe_joint,
      args: List[Any],
      *,
      is_joint_structure: bool,
      meta: ViewAndMutationMeta,
      fw_only: Callable,
  ) -> SubclassTracingInfo:
      # Skip logic if we don't need to trace through any subclasses
      req_subclass_dispatch = requires_subclass_dispatch(args, meta)
      if not req_subclass_dispatch:
          return SubclassTracingInfo(
              plain_tensor_trace_fn=flat_fn_maybe_joint,
              plain_tensor_args=args,
              maybe_subclass_meta=None,
          )
  
      # TODO: add subclass guards (later PR).
  
      # What's going on here? We need to compute subclass metadata about the outputs of the joint (grad_inputs).
      # Annoying: we don't know the grad input metas until we're in the middle of tracing the joint,
      # so we set it later, while we're tracing the joint (see inner_fn() below).
      # Another option would be to run our run_functionalized_fw_and_collect_metadata() function
      # directly on the joint, but this would hurt compile time (adding yet another pass through the joint).
      subclass_meta = SubclassMeta()
  
      def inner_fn(fn, args, *, use_trace_joint: bool):
          # Step 1: wrap tensor inputs into subclasses if necessary
          all_args = wrap_tensor_subclasses_maybe_joint(
              args, is_joint_structure=use_trace_joint, meta=meta
          )
  
          # Step 2: call the inner function, with our (maybe subclass) inputs
          wrapped_outs = fn(*all_args)
  
          if use_trace_joint:
              # See Note: [Computing Subclass Metadata about grad_inputs]
              # We also stash subclass info on our grad_inputs, if we're tracing the joint.
              nonlocal subclass_meta
              assert (
                  isinstance(wrapped_outs, tuple)
                  and len(wrapped_outs) == 2
              )
              # Don't need fw outs since we already have subclass metadata on them
              grad_inputs = wrapped_outs[1]
              subclass_meta.grad_input_metas = create_subclass_meta(
                  grad_inputs
              )
  
          # Step 3: Unwrap any subclass outputs back into dense tensors
          unwrapped_outs = unwrap_tensor_subclasses(
              wrapped_outs, is_joint_structure=use_trace_joint
          )
          return unwrapped_outs
  
      def joint_fn(primals, tangents):
          return inner_fn(
              flat_fn_maybe_joint,
              (primals, tangents),
              use_trace_joint=True,
          )
  
      def fw_fn(*primals):
          return inner_fn(
              flat_fn_maybe_joint, primals, use_trace_joint=False
          )
  
      def metadata_fn(*primals):
          return inner_fn(fw_only, primals, use_trace_joint=False)
  
      args_unwrapped = unwrap_tensor_subclasses(
          args, is_joint_structure=is_joint_structure
      )
  
      if is_joint_structure:
          primals_unwrapped = args_unwrapped[0]
          fn_to_trace = joint_fn
      else:
          primals_unwrapped = args_unwrapped
          fn_to_trace = fw_fn
  
      # Note: [Partitioner handling for Subclasses, Part 1]
      # The way the partitioner works is that:
      # (1) we pass is a single graph containing the joint fw/bw,
      #     where the # of graph outputs corresponds to # fw_outputs + # grad_inputs
      # (2) The partitioner accepts an arguments, num_fwd_outputs,
      #     and assumes that the first "num_fwd_outputs" graph outputs correspond
      #     to outputs of the forward graph.
      # How do tensor subclasses enter the picture?
      # the num_fwd_outputs in the final graph is actually non-trivial to compute,
      # because it can be influenced by input mutations and intermediate bases.
      # So we compute it by inspecting the current ViewAndMutationMeta object.
      # However, the original ViewAndMutationMeta that we computed was created
      # on the subclass -> subclass graph,
      # which can have a different number of outputs than the dense -> dense graph.
      # That's why we createa a fresh metadata object on the dense -> dense function here,
      # and plumb it back up to the partitioner.
      # See Note: [Partitioner handling for Subclasses, Part 2] for more info.
      meta_updated = run_functionalized_fw_and_collect_metadata(
          metadata_fn,
          keep_input_mutations=meta.keep_input_mutations,
          is_train=meta.is_train,
          requires_subclass_dispatch=True,
      )(*primals_unwrapped)
  
      subclass_meta.fw_metadata = meta_updated
  
      return SubclassTracingInfo(
          plain_tensor_trace_fn=fn_to_trace,
          plain_tensor_args=args_unwrapped,
          maybe_subclass_meta=subclass_meta,
      )
  
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 6 for calling a global function with example usage: aot_dispatch_subclass(flat_fn_maybe_joint, args) and returns: SubclassTracingInfo, SubclassTracingInfo, unwrapped_outs, inner_fn, inner_fn, inner_fn