code: def _run_sharded_embedding_bag(
      self,
      spec,
      input_size,
      num_embeddings,
      embedding_dim,
      mode,
      include_last_offset=False,
      offset_size=None,
      max_norm=None,
      norm_type=2.0,
      padding_idx=None,
  ):
      # Use same seed.
      torch.manual_seed(0)
      local_embedding_bag = torch.nn.EmbeddingBag(
          num_embeddings,
          embedding_dim,
          mode=mode,
          max_norm=max_norm,
          norm_type=norm_type,
          include_last_offset=include_last_offset,
          padding_idx=padding_idx,
      ).cuda(self.rank)
  
      sharded_embedding_bag = torch.nn.EmbeddingBag(
          num_embeddings,
          embedding_dim,
          mode=mode,
          max_norm=max_norm,
          norm_type=norm_type,
          include_last_offset=include_last_offset,
          padding_idx=padding_idx,
      )
  
      # Copy the weights from local embedding bag.
      sharded_embedding_bag.weight = clone_module_parameter(
          local_embedding_bag, "weight"
      )
  
      # Shard the parameter.
      shard_parameter(sharded_embedding_bag, "weight", spec)
  
      # Run sharded computation
      torch.manual_seed(self.rank)  # inputs different on each rank
      inp = torch.randint(0, num_embeddings, tuple(input_size)).cuda(
          self.rank
      )
      per_sample_weights = None
      if mode == "sum":
          per_sample_weights = torch.rand(*input_size).cuda(self.rank)
  
      offsets = None
      if len(input_size) == 1:
          # We need to generate certain length offset for each rank.
          # The current implementation and dist API does not support
          # the case when the offset has different lengths.
          # input_size[0] >> offset_size, so the while loop will not
          # for too long.
          while offsets is None or (offsets.size(0) != offset_size):
              offsets = torch.randint(input_size[0], (offset_size,))
              offsets[0] = 0
              if include_last_offset:
                  offsets[-1] = input_size[0]
              offsets = (
                  torch.unique(offsets, sorted=True)
                  .contiguous()
                  .cuda(self.rank)
              )
  
      # If max_norm is set, we need to ensure that the renorm has been applied across
      # inputs from all ranks.
      if max_norm is not None:
          gathered_inputs = [
              torch.zeros_like(inp) for _ in range(TEST_GPU_NUM)
          ]
          dist.all_gather(gathered_inputs, inp)
          unique_inp = torch.unique(torch.cat(gathered_inputs))
          offsets_dummy = torch.tensor([len(unique_inp) // 2]).cuda(
              self.rank
          )
          local_embedding_bag(unique_inp, offsets=offsets_dummy)
  
      sharded_output = sharded_embedding_bag(
          inp,
          offsets=offsets,
          per_sample_weights=per_sample_weights,
      )
  
      # Run local computation
      local_output = local_embedding_bag(
          inp,
          offsets=offsets,
          per_sample_weights=per_sample_weights,
      )
  
      # Compare local weight and shared one to ensure the renorm
      # as expected.
      if max_norm is not None:
          sharded_dim = spec.dim
          sharded_weight = sharded_embedding_bag.weight.local_shards()[
              0
          ].tensor
          (
              start_pos,
              chunk_size,
          ) = generate_local_weight_sharding_params_for_test(
              local_embedding_bag.weight,
              sharded_dim,
              TEST_GPU_NUM,
              spec,
              self.rank,
          )
          local_weight_narrowed = local_embedding_bag.weight.narrow(
              sharded_dim, start_pos, chunk_size
          )
          self.assertEqual(local_weight_narrowed, sharded_weight)
  
      # Verify
      self.assertEqual(local_output, sharded_output)
  
      # Validate for torch.nn.functional.embedding_bag version.
      local_output = torch.nn.functional.embedding_bag(
          inp,
          local_embedding_bag.weight,
          offsets=offsets,
          mode=mode,
          per_sample_weights=per_sample_weights,
          include_last_offset=include_last_offset,
          max_norm=max_norm,
          norm_type=norm_type,
          padding_idx=padding_idx,
      )
      sharded_output = torch.nn.functional.embedding_bag(
          inp,
          sharded_embedding_bag.weight,
          offsets=offsets,
          mode=mode,
          per_sample_weights=per_sample_weights,
          include_last_offset=include_last_offset,
          max_norm=max_norm,
          norm_type=norm_type,
          padding_idx=padding_idx,
      )
  
      self.assertEqual(local_output, sharded_output)
  
file_path: pytorch/test/distributed/_shard/sharded_tensor/ops/test_embedding_bag.py
imports: import sys
  import torch
  import torch.distributed
  from torch.distributed._shard import shard_parameter
  from torch.testing._internal.common_distributed import requires_nccl, skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import TEST_WITH_DEV_DBG_ASAN, run_tests
  from torch.testing._internal.distributed._shard.sharded_tensor import TEST_GPU_NUM, ShardedTensorTestBase, with_comms
  from torch.testing._internal.distributed._shard.sharded_tensor._test_ops_common import clone_module_parameter, generate_chunk_sharding_specs_for_test, generate_local_weight_sharding_params_for_test
label_desc: Use image node_id 1 for calling the TestShardedEmbeddingBag obj's underlying member method code with example usage: obj._run_sharded_embedding_bag(spec, input_size, num_embeddings, embedding_dim, mode, include_last_offset, offset_size, max_norm, norm_type, padding_idx) without return types