code: def parallel_lm_logits(
      input_: torch.Tensor,
      word_embeddings_weight: torch.Tensor,
      parallel_output: bool,
      bias: torch.Tensor = None,
      async_tensor_model_parallel_allreduce: bool = False,
      sequence_parallel: bool = False,
      gradient_accumulation_fusion: bool = False,
  ):
      """Language Model logits using word embedding weights.
  
      Args:
          input_ (torch.Tensor): [b, s, h]
          word_embeddings_weight (torch.Tensor): [(padded) vocab size, h]
          parallel_output (bool): False will gather logits from tensor model parallel region
          bias (torch.Tensor, optional): bias tensor. Defaults to None.
          async_tensor_model_parallel_allreduce (bool, optional): Defaults to False.
          sequence_parallel (bool, optional): If True will use sequence parallelism. Defaults to False.
          gradient_accumulation_fusioa (bool, optional): If True fuse gradient accumulation to WGRAD GEMM
  
      Returns:
          torch.Tensor: [b, s, (padded) vocab size]
      """
  
      tensor_model_parallel = (
          parallel_state.get_tensor_model_parallel_world_size() > 1
      )
  
      # async grad allreduce can only be used when not using sequence parallelism
      async_grad_allreduce = (
          async_tensor_model_parallel_allreduce
          and tensor_model_parallel
          and not sequence_parallel
      )
  
      # copy input_ to model parallel region if needed
      if async_tensor_model_parallel_allreduce or sequence_parallel:
          input_parallel = input_
  
      else:
          input_parallel = (
              tensor_parallel.copy_to_tensor_model_parallel_region(
                  input_
              )
          )
  
      # Matrix multiply.
      logits_parallel = (
          linear_with_grad_accumulation_and_async_allreduce(
              input=input_parallel,
              weight=word_embeddings_weight,
              bias=bias,
              gradient_accumulation_fusion=gradient_accumulation_fusion,
              async_grad_allreduce=async_grad_allreduce,
              sequence_parallel=sequence_parallel,
          )
      )
  
      # Gather if needed.
      if parallel_output:
          return logits_parallel
      else:
          return (
              tensor_parallel.gather_from_tensor_model_parallel_region(
                  logits_parallel
              )
          )
  
imports: import itertools
  import math
  from typing import Dict, Iterator, List, Tuple, Union
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 2 for calling a global function with example usage: parallel_lm_logits(input_, word_embeddings_weight, parallel_output, bias, async_tensor_model_parallel_allreduce, sequence_parallel, gradient_accumulation_fusion) and returns: logits_parallel, tensor_parallel