code: def test_cross_entropy(self):
      batch_size, sequence_length, vocab_size_per_partition = 13, 17, 11
      logits_scale = 1000.0
      seed = 1234
      for tensor_model_parallel_world_size in range(
          1, self.world_size + 1
      ):
          if self.world_size % tensor_model_parallel_world_size:
              continue
          parallel_state.initialize_model_parallel(
              tensor_model_parallel_size_=tensor_model_parallel_world_size,
          )
          vocab_size = (
              vocab_size_per_partition
              * tensor_model_parallel_world_size
          )
          loss_torch, grad_torch = torch_cross_entropy(
              batch_size,
              sequence_length,
              vocab_size,
              logits_scale,
              seed,
          )
          (
              loss_tensor_parallel,
              grad_tensor_parallel,
          ) = tensor_sharded_cross_entropy(
              batch_size,
              sequence_length,
              vocab_size,
              logits_scale,
              seed,
          )
  
          self.assertEqual(
              loss_torch,
              loss_tensor_parallel,
              msg=f"tensor_model_parallel_size: {tensor_model_parallel_world_size}",
          )
          self.assertEqual(
              grad_torch,
              grad_tensor_parallel,
              msg=f"tensor_model_parallel_size: {tensor_model_parallel_world_size}",
          )
  
          parallel_state.destroy_model_parallel()
  
imports: import logging
  from typing import Tuple
  import torch
  import torch.nn.functional
  from torch.testing._internal import common_utils
  from apex.transformer import parallel_state
  from apex.transformer import tensor_parallel
  from apex.transformer.tensor_parallel import cross_entropy
  from apex.transformer.testing.commons import set_random_seed, IdentityLayer
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
  from apex.transformer.testing.distributed_test_base import UccDistributedTestBase
label_desc: Use image node_id 1 for calling the VocabParallelCrossEntropyTestBase obj's underlying member method code with example usage: obj.test_cross_entropy() without return types