code: def _initialize_affine_weight_gpu(
      weight,
      init_method,
      partition_dim,
      stride=1,
      expert_parallel=False,
  ):
      """Initialize affine weight for model parallel on GPU."""
  
      set_tensor_model_parallel_attributes(
          tensor=weight,
          is_parallel=True,
          dim=partition_dim,
          stride=stride,
      )
  
      if not expert_parallel:
          with get_cuda_rng_tracker().fork():
              init_method(weight)
      else:
          with get_cuda_rng_tracker().fork(
              get_expert_parallel_rng_tracker_name()
          ):
              init_method(weight)
  
imports: import math
  import os
  import warnings
  from typing import Callable, Optional
  import torch
  import torch.nn.functional
  import torch.nn.init
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.nn.parameter import Parameter
  from megatron.core.model_parallel_config import ModelParallelConfig
  from megatron.core.parallel_state import get_global_memory_buffer, get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
  from .mappings import copy_to_tensor_model_parallel_region, gather_from_sequence_parallel_region, gather_from_tensor_model_parallel_region, reduce_from_tensor_model_parallel_region, reduce_scatter_to_sequence_parallel_region, scatter_to_tensor_model_parallel_region
  from .random import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
  from .utils import VocabUtility, divide, split_tensor_along_last_dim
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, expert_parallel) without return types