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 = scatter_to_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: fairscale/tests/nn/model_parallel/test_cross_entropy.py
imports: import torch
  import torch.nn.functional
  from fairscale.fair_dev.testing.testing import IdentityLayer, dist_init, set_random_seed, spawn_for_all_world_sizes
  from fairscale.nn.model_parallel import initialize
  from fairscale.nn.model_parallel.cross_entropy import vocab_parallel_cross_entropy
  from fairscale.nn.model_parallel.mappings import scatter_to_model_parallel_region
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