code: def _initialize_affine_weight_gpu(
      weight, init_method, partition_dim, stride=1
  ):
      """Initialize affine weight for model parallel on GPU.
  
      Args:
          weight (Parameter):
          init_method (Callable[[Tensor], None]): Taking a Tensor and initialize its elements.
          partition_dim (int): Dimension to apply partition.
          stride (int):
      """
  
      set_tensor_model_parallel_attributes(
          tensor=weight,
          is_parallel=True,
          dim=partition_dim,
          stride=stride,
      )
  
      with get_cuda_rng_tracker().fork():
          init_method(weight)
  
imports: from typing import Optional, Dict, Tuple, List
  import warnings
  import torch
  import torch.nn.functional
  import torch.nn.init
  from torch.nn.parameter import Parameter
  from apex._autocast_utils import _cast_if_autocast_enabled
  from apex.transformer.parallel_state import get_tensor_model_parallel_group
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
  from apex.transformer.utils import divide
  from apex.transformer.tensor_parallel.mappings import copy_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import scatter_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_scatter_to_sequence_parallel_region
  from apex.transformer.tensor_parallel.random import get_cuda_rng_tracker
  from apex.transformer.tensor_parallel.utils import VocabUtility
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 5 for calling a global function with example usage: _initialize_affine_weight_gpu(weight, init_method, partition_dim, stride) without return types