code: def mpu_cross_entropy(
      batch_size, seq_length, vocab_size, logits_scale, seed
  ):
      set_random_seed(seed)
      identity = IdentityLayer(
          (batch_size, seq_length, vocab_size), scale=logits_scale
      ).cuda()
      logits = identity()
      logits_parallel = mpu.scatter_to_tensor_model_parallel_region(
          logits
      )
      target = torch.cuda.LongTensor(
          size=(batch_size, seq_length)
      ).random_(0, vocab_size)
      loss = vocab_parallel_cross_entropy(
          logits_parallel, target
      ).mean()
      loss.backward()
      return loss, identity.weight.grad
  
file_path: megatron-lm/megatron/mpu/tests/test_cross_entropy.py
imports: from commons import set_random_seed
  from commons import IdentityLayer
  from commons import print_separator
  from commons import initialize_distributed
  from mpu.cross_entropy import vocab_parallel_cross_entropy
  import mpu
  import torch.nn.functional
  import torch
  import random
  import sys
label_desc: Use image node_id 2 for calling a global function with example usage: mpu_cross_entropy(batch_size, seq_length, vocab_size, logits_scale, seed) and returns: loss,  identity