code: def _get_extra_te_kwargs(config: TransformerConfig):
      extra_transformer_engine_kwargs = {
          "params_dtype": config.params_dtype,
      }
  
      te_version = packaging.version.Version(
          version("transformer-engine")
      )
      if te_version >= packaging.version.Version("0.12.0"):
          if config.use_cpu_initialization:
              extra_transformer_engine_kwargs["device"] = "cpu"
          else:
              extra_transformer_engine_kwargs[
                  "device"
              ] = torch.cuda.current_device()
      return extra_transformer_engine_kwargs
  
imports: import os
  from importlib.metadata import version
  from typing import Callable
  import torch
  import transformer_engine
  from pkg_resources import packaging
  from torch import Tensor
  from megatron.core import ModelParallelConfig
  from megatron.core.parallel_state import get_context_parallel_global_ranks, get_context_parallel_group, get_tensor_model_parallel_group
  from megatron.core.tensor_parallel import get_cuda_rng_tracker
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
label_desc: Use image node_id 1 for calling a global function with example usage: _get_extra_te_kwargs(config) and returns: extra_transformer_engine_kwargs