code: def ddp_ref():
      # Cover different bias flavors. Use random instead of parameterize them to reduce
      # the test runtime. Otherwise, we would have covered all cases exhaustively.
      conv_bias = True
      if random.randint(0, 1) == 0:
          conv_bias = False
      linear_bias = True
      if random.randint(0, 1) == 0:
          linear_bias = False
  
      # Get a reference model state
      model = Model(conv_bias, linear_bias)
      state_before = model.state_dict()
  
      # Get reference inputs per rank.
      world_size = _world_size
      iterations = _iterations
      print(
          f"Getting DDP reference for world_size {world_size} and iterations {iterations}"
      )
      inputs = [[] for i in range(world_size)]
      for rank in range(world_size):
          for i in range(iterations):
              inputs[rank].append(torch.rand(2, 2, 2, 2))
  
      # Run reference DDP training 4 times, fp and mp, sync_bn or not.
      state_after = {}
      for precision, sync_bn in product(
          ["full", "mixed"], ["none", "pytorch"]
      ):
          temp_file_name = tempfile.mkstemp()[1]
          unused = tempfile.mkstemp()[1]
          rank_0_output = tempfile.mkstemp()[1]
          try:
              fsdp_config = (
                  None  # This means we use DDP in _distributed_worker.
              )
              mp.spawn(
                  _distributed_worker,
                  args=(
                      world_size,
                      fsdp_config,
                      None,
                      precision == "mixed",
                      temp_file_name,
                      unused,
                      state_before,
                      inputs,
                      rank_0_output,
                      None,
                      sync_bn,
                      conv_bias,
                      linear_bias,
                  ),
                  nprocs=world_size,
                  join=True,
              )
              state_after[(precision, sync_bn)] = torch.load(
                  rank_0_output
              )
          finally:
              rmf(temp_file_name)
              rmf(unused)
              rmf(rank_0_output)
  
      # Sanity check DDP's final states.
      states = list(state_after.values())
      for state in states[1:]:
          assert state_dict_norm(states[0]) != state_dict_norm(state)
  
      return state_before, inputs, conv_bias, linear_bias, state_after
  
imports: import contextlib
  from itertools import product
  import random
  import tempfile
  import pytest
  import torch
  from torch.cuda.amp import GradScaler
  import torch.multiprocessing
  from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, CrossEntropyLoss, Linear, Module, ReLU, Sequential, Sigmoid, SyncBatchNorm
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim import SGD
  from fairscale.fair_dev.testing.testing import dist_init, objects_are_equal, rmf, skip_if_single_gpu, state_dict_norm, teardown, torch_cuda_version
  from fairscale.internal import torch_version
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import TrainingState, auto_wrap_bn
label_desc: Use image node_id 1 for calling a global function with example usage: ddp_ref() and returns: state_before,  inputs,  conv_bias,  linear_bias,  state_after