class_docstr: Tests ``fully_shard`` initialization.
class_docstr_tok: ["Tests","``","fully_shard","``","initialization","."]
code: def _test_nested_fully_shard_shared_state(self, use_policy: bool):
      device = torch.device("cuda")
      composable_module = CompositeParamModel(device=device)
      if use_policy:
          fully_shard(
              composable_module, policy=ModuleWrapPolicy({UnitModule})
          )
      else:
          fully_shard(composable_module.u1)
          fully_shard(composable_module.u2)
          fully_shard(composable_module)
  
      # Run a forward pass to trigger lazy initialization
      inp = torch.randn((2, 100), device=device)
      composable_module(inp)
  
      # Check that all modules with `fully_shard` applied share the same data
      # structure state for the structures with the given names (there is no
      # need to check all of them to verify that the sharing worked).
      # NOTE: This check only requires that the data structure state is
      # shared. Namely, sharing the FSDP state object itself is sufficient
      # but not necessary.
      data_structure_names = [
          "_exec_order_data",
          "_free_event_queue",
          "_pre_unshard_stream",
          "_unshard_stream",
          "_post_backward_stream",
          "_default_stream",
      ]
      for data_structure_name in data_structure_names:
          all_structures = set()
          for module in (
              composable_module.u1,
              composable_module.u2,
              composable_module,
          ):
              all_structures.add(
                  id(
                      getattr(
                          fully_shard.state(module), data_structure_name
                      )
                  )
              )
          self.assertEqual(len(all_structures), 1)
  
imports: import copy
  import sys
  from typing import Optional
  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 BackwardPrefetch, FullyShardedDataParallel
  from torch.distributed.fsdp._common_utils import _is_fsdp_flattened, clean_tensor_name
  from torch.distributed.fsdp.wrap import _Policy, CustomPolicy, ModuleWrapPolicy
  from torch.testing._internal.common_dist_composable import CompositeParamModel, FakeSequential, NestedSequentialModel, 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 10 for calling the TestInitialization obj's underlying member method code with example usage: obj._test_nested_fully_shard_shared_state(use_policy) without return types