code: def tensor_sharded_cross_entropy(
      batch_size,
      seq_length,
      vocab_size,
      logits_scale,
      seed,
      label_smoothing=0.0,
  ):
      set_random_seed(seed)
      identity = IdentityLayer(
          (batch_size, seq_length, vocab_size), scale=logits_scale
      ).cuda()
      logits = identity()
      logits_parallel = (
          tensor_parallel.scatter_to_tensor_model_parallel_region(
              logits
          )
      )
      target = torch.cuda.LongTensor(
          size=(batch_size, seq_length)
      ).random_(0, vocab_size)
      logits_parallel_ = logits_parallel.clone().detach()
      loss = cross_entropy.vocab_parallel_cross_entropy(
          logits_parallel, target, label_smoothing=label_smoothing
      ).mean()
      loss.backward()
      # check for mutation
      assert torch.equal(logits_parallel_, logits_parallel)
      return loss, identity.weight.grad
  
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 2 for calling a global function with example usage: tensor_sharded_cross_entropy(batch_size, seq_length, vocab_size, logits_scale, seed, label_smoothing) and returns: loss,  identity