class_docstr: Tests ``fully_shard`` runtime (forward/backward/optimizer).
code: def _test_unshard_reshard_order(self, fsdp_wrap_mode: FSDPWrapMode):
      device = torch.device("cuda")
      (
          composable_module,
          composable_optim,
          fsdp_wrapped_model,
          fsdp_wrapped_optim,
      ) = self._init_models_and_optims(
          device, fsdp_wrap_mode, ShardingStrategy.FULL_SHARD
      )
      # Before checking the unshard/reshard order, sanity check that the
      # assumption about wrapper FQN being a suffix of composable FQN holds
      all_composable_handles = traversal_utils._get_fsdp_handles(
          composable_module
      )
      all_wrapped_handles = traversal_utils._get_fsdp_handles(
          fsdp_wrapped_model
      )
      for c_handle, w_handle in zip(
          all_composable_handles, all_wrapped_handles
      ):
          self._check_same_param_handles(c_handle, w_handle)
      num_handles = len(all_composable_handles)
  
      orig_unshard = torch.distributed.fsdp._runtime_utils._unshard
      orig_reshard = torch.distributed.fsdp._runtime_utils._reshard
      UnshardReshardEvent = Tuple[str, FlatParamHandle]
  
      def patched_unshard(
          unshard_reshard_order: List[UnshardReshardEvent],
          state: _FSDPState,
          handle: FlatParamHandle,
          *args,
          **kwargs,
      ):
          unshard_reshard_order.append(("unshard", handle))
          return orig_unshard(state, handle, *args, **kwargs)
  
      def patched_reshard(
          unshard_reshard_order: List[UnshardReshardEvent],
          state: _FSDPState,
          handle: FlatParamHandle,
          *args,
          **kwargs,
      ):
          unshard_reshard_order.append(("reshard", handle))
          return orig_reshard(state, handle, *args, **kwargs)
  
      @contextlib.contextmanager
      def patch_unshard(_patched_unshard: Callable):
          _orig_unshard = torch.distributed.fsdp._runtime_utils._unshard
          torch.distributed.fsdp._runtime_utils._unshard = (
              _patched_unshard
          )
          try:
              yield
          finally:
              torch.distributed.fsdp._runtime_utils._unshard = (
                  _orig_unshard
              )
  
      @contextlib.contextmanager
      def patch_reshard(_patched_reshard: Callable):
          _orig_reshard = torch.distributed.fsdp._runtime_utils._reshard
          torch.distributed.fsdp._runtime_utils._reshard = (
              _patched_reshard
          )
          try:
              yield
          finally:
              torch.distributed.fsdp._runtime_utils._unshard = (
                  _orig_reshard
              )
  
      composable_order: List[UnshardReshardEvent] = []
      wrapped_order: List[UnshardReshardEvent] = []
  
      inp = torch.randn(2, 100, device="cuda")
      losses: List[torch.Tensor] = []
  
      for order, model, optim in (
          (composable_order, composable_module, composable_optim),
          (wrapped_order, fsdp_wrapped_model, fsdp_wrapped_optim),
      ):
          with patch_unshard(
              functools.partial(patched_unshard, order)
          ), patch_reshard(functools.partial(patched_reshard, order)):
              optim.zero_grad(set_to_none=True)
              out = model(inp)
              loss = out.sum()
              losses.append(loss)
              loss.backward()
              optim.step()
      self.assertEqual(losses[0], losses[1])
  
      # Sanity check that the unshard/reshard events were recorded, where we
      # expect one unshard/reshard pair for forward, one pair for backward,
      # and possibly some extra unshards from backward prefetching (in this
      # case, we expect exactly 2 extra since there are 3 handles)
      self.assertGreaterEqual(
          len(composable_order), 2 * 2 * num_handles
      )
      self.assertGreaterEqual(len(wrapped_order), 2 * 2 * num_handles)
      self.assertGreaterEqual(
          len([e for e in composable_order if e[0] == "unshard"]),
          2 * num_handles,
      )
      self.assertGreaterEqual(
          len([e for e in wrapped_order if e[0] == "unshard"]),
          2 * num_handles,
      )
      self.assertGreaterEqual(
          len([e for e in composable_order if e[0] == "reshard"]),
          2 * num_handles,
      )
      self.assertGreaterEqual(
          len([e for e in wrapped_order if e[0] == "reshard"]),
          2 * num_handles,
      )
  
      # Check that the unshard/reshard order matches
      self.assertEqual(len(composable_order), len(wrapped_order))
      for (
          (composable_event, composable_handles_key),
          (wrapped_event, wrapped_handles_key),
      ) in zip(composable_order, wrapped_order):
          self.assertEqual(composable_event, wrapped_event)
          self._check_same_param_handles(
              composable_handles_key, wrapped_handles_key
          )
  
imports: import contextlib
  import copy
  import functools
  import sys
  from enum import auto, Enum
  from typing import Callable, List, Tuple
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed._composable import fully_shard
  from torch.distributed.fsdp import FullyShardedDataParallel, ShardingStrategy
  from torch.distributed.fsdp._common_utils import _FSDPState
  from torch.distributed.fsdp._flat_param import FlatParamHandle
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.testing._internal.common_dist_composable import CompositeParamModel, UnitModule
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 6 for calling the TestRuntime obj's underlying member method code with example usage: obj._test_unshard_reshard_order(fsdp_wrap_mode) and returns: orig_unshard, orig_reshard