code: def test_cross_entropy(tensor_model_parallel_size):
      if torch.distributed.get_rank() == 0:
          print(
              "> testing cross entropy with model parallel size {} ...".format(
                  tensor_model_parallel_size
              )
          )
  
      mpu.initialize_model_parallel(tensor_model_parallel_size)
      tensor_model_parallel_size = (
          mpu.get_tensor_model_parallel_world_size()
      )
  
      batch_size = 13
      seq_length = 17
      vocab_size_per_partition = 11
      logits_scale = 1000.0
      vocab_size = vocab_size_per_partition * tensor_model_parallel_size
      seed = 1234
  
      loss_torch, grad_torch = torch_cross_entropy(
          batch_size, seq_length, vocab_size, logits_scale, seed
      )
      loss_mpu, grad_mpu = mpu_cross_entropy(
          batch_size, seq_length, vocab_size, logits_scale, seed
      )
  
      error = loss_torch.sub_(loss_mpu).abs().max()
      print(
          "   max error in loss on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6
  
      error = grad_torch.sub_(grad_mpu).abs().max()
      print(
          "   max error in grad on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6
  
      # Reset groups
      mpu.destroy_tensor_model_parallel()
  
      torch.distributed.barrier()
      if torch.distributed.get_rank() == 0:
          print(">> passed the test :-)")
  
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 3 for calling a global function with example usage: test_cross_entropy(tensor_model_parallel_size) without return types