docstr: Tests ``writeback=True`` for a singleton parameter, which includes
  testing that writing to padding does not persist.
  NOTE: This method depends on FSDP internals.
docstr_tok: ["Tests","``","writeback=True","``","for","a","singleton","parameter",",","which","includes","testing","that","writing","to","padding","does","not","persist",".","NOTE",":","This","method","depends","on","FSDP","internals","."]
code: def test_unshard_singleton_param_writeback(self):
      """
      Tests ``writeback=True`` for a singleton parameter, which includes
      testing that writing to padding does not persist.
      NOTE: This method depends on FSDP internals.
      """
      model = FSDP(nn.Linear(1, 1, bias=False, device=self.device))
      flat_param = model._handle.flat_param
      self.assertEqual(1, flat_param.numel())
      # Write a known value to the *sharded* `FlatParameter`
      with torch.no_grad():
          # For nonzero ranks, this write is to padding
          flat_param[0] = self.rank + 2
      with FSDP.summon_full_params(model, writeback=True):
          self.assertEqual(1, flat_param.numel())
          with torch.no_grad():
              flat_param.zero_()
      # NOTE: This checks that writes to padding did not persist, which is
      # *not* strictly required for correctness.
      if self.rank == 0:  # did not write to padding
          self.assertEqual(0, flat_param[0])
      else:  # wrote to padding
          self.assertEqual(self.rank + 2, flat_param[0])
  
imports: import contextlib
  import itertools
  import math
  import sys
  from typing import Any, Dict, List, Optional, Union
  import torch
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch import distributed
  from torch.distributed.fsdp import CPUOffload, FullyShardedDataParallel, MixedPrecision, ShardingStrategy
  from torch.distributed.fsdp._common_utils import clean_tensor_name
  from torch.distributed.fsdp._flat_param import FlatParameter
  from torch.nn.parallel.distributed import DistributedDataParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import CUDAInitMode, FSDPInitMode, FSDPTest, NestedWrappedModule, TransformerWithSharedParams
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 4 for calling the TestUnshardParams obj's underlying member method code with example usage: obj.test_unshard_singleton_param_writeback() without return types