code: def _test_basic_func(
      rank, ddp_cls, world_size, tempfile_name, test_case
  ):
      _dist_init(
          rank, world_size, tempfile_name, backend="nccl"
      )  # Covers nccl
  
      model = Linear(2, 2)
      model.to("cuda")
      if ddp_cls is DDP:
          model = ddp_cls(model, device_ids=[rank])
          optim = AdaScale(SGD(model.parameters(), lr=0.1))
      elif ddp_cls is SDP:
          optim = AdaScale(OSS(model.parameters(), SGD, lr=0.1))
          model = ddp_cls(model, sharded_optimizer=optim)
      else:
          assert ddp_cls is FSDP, ddp_cls
          # Two cases:
          #    flatten=True : AdaScale wrapper must be after FSDP and it receives
          #                   a single grad tensor. It won't receive grad if
          #                   wrapped before.
          #    flatten=False: AdaScale can be both before or after FSDP.
          # So, it is better to do AdaScale after FSDP.
          model = ddp_cls(model, flatten_parameters=False)
          optim = AdaScale(SGD(model.parameters(), lr=0.1))
      if "input" in test_case:
          # single iter
          in_data = Tensor(test_case["input"][rank])
          in_data = in_data.cuda()
          out = model(in_data)
          out.sum().backward()
          if ddp_cls is DDP:
              assert np.allclose(
                  optim.gain(), test_case["expected_gain"]
              ), optim.gain()
              w, b = model.parameters()
              assert np.allclose(
                  w.grad.cpu(), test_case["expected_grad"]
              ), w.grad
              assert np.allclose(
                  b.grad.cpu(), test_case["expected_bias_grad"]
              ), b.grad
          optim.step()
          optim.zero_grad()
      else:
          # multiple iters
          n = len(test_case["inputs"])
          for i, in_data in enumerate(test_case["inputs"]):
              in_data = Tensor(in_data[rank]).cuda()
              out = model(in_data)
              out.sum().backward()
              if i == n - 1 and ddp_cls is DDP:
                  assert np.allclose(
                      optim.gain(), test_case["expected_gain"]
                  ), optim.gain()
                  w, b = model.parameters()
                  assert np.allclose(
                      w.grad.cpu(), test_case["expected_grad"]
                  ), w.grad
                  assert np.allclose(
                      b.grad.cpu(), test_case["expected_bias_grad"]
                  ), b.grad
              optim.step()
              optim.zero_grad()
  
      dist.destroy_process_group()
  
imports: import tempfile
  import numpy
  import pytest
  import torch
  from torch import Tensor
  import torch.distributed
  import torch.multiprocessing
  from torch.nn import Linear
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim import SGD
  from fairscale.fair_dev.testing.golden_testing_data import adascale_test_data
  from fairscale.fair_dev.testing.testing import skip_if_single_gpu
  from fairscale.internal import torch_version
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import ShardedDataParallel
  from fairscale.optim import OSS, AdaScale
label_desc: Use image node_id 2 for calling a global function with example usage: _test_basic_func(rank, ddp_cls, world_size, tempfile_name, test_case) without return types