code: def _handle_col_wise_sharding(
      input,
      world_size,
      weight,
      local_shard,
      max_norm,
      norm_type,
      padding_idx,
      pg,
  ):
      """
      Entry-point function to handle the logic of col-wise sharding of weight
      for embedding. (Detailed explanations of the logic can be found in
      the comment for sharded_embedding.)
  
      Args:
          input: list of ID used for lookup and aggregation.
          world_size: number of ranks.
          weight: sharded weight tensor.
          local_shard: col-wise shared local weight used for lookup.
          max_norm: If given, each embedding vector with norm larger
              than max_norm is renormalized to have norm max_norm.
              Note: this will modify weight in-place.
          norm_type: The p in the p-norm to compute for the max_norm option.
          padding_idx: If specified, the entries at padding_idx do
              not contribute to the gradient; therefore, the embedding
              vector at padding_idx is not updated during training,
              i.e. it remains as a fixed “pad”.
          pg: process group.
  
      Returns: final result of lookup.
      """
      # allgather the inputs first for non Replicated Tensor.
      gathered_inputs = all_gather(input, group=pg)
  
      if max_norm is not None:
          # max_norm changes the weight in-place
          local_shard = _handle_max_norm_col_wise(
              max_norm,
              norm_type,
              local_shard,
              input,
              world_size,
              gathered_inputs,
              pg,
          )
  
      output = _handle_col_wise_sharding_base(
          torch.nn.functional.embedding,
          len(input.size()),
          input,
          world_size,
          weight,
          local_shard,
          pg,
          gathered_inputs,
          padding_idx=padding_idx,
      )
      return (output, local_shard)
  
file_path: pytorch/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding.py
imports: import torch
  import torch.distributed
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._shard.sharding_spec import ChunkShardingSpec
  from torch.distributed._shard.sharding_spec.api import custom_sharding_spec_op
  from torch.distributed.nn.functional import all_gather, reduce_scatter
  from ._common import _all_gather_base_input, _handle_col_wise_sharding_base, _handle_max_norm_col_wise, _handle_row_wise_mask
label_desc: Use image node_id 3 for calling a global function with example usage: _handle_col_wise_sharding(input, world_size, weight, local_shard, max_norm, norm_type, padding_idx, pg) and returns: output,  local_shard