code: def test_parallel_self_attention(tensor_model_parallel_size):
      if torch.distributed.get_rank() == 0:
          print(
              "> testing ParallelSelfAttention with model parallel "
              "size: {}".format(tensor_model_parallel_size)
          )
  
      num_att_heads_per_partition = 3
      hidden_size_per_att_head = 7
      dropout_prob = 0.0  # has to be zero
      batch_size = 5
      sequence_length = 13
  
      (
          rank_1,
          hideen_size_1,
          tensor_model_parallel_size_1,
          loss_1,
          attention_layer_1,
          identity_layer_1,
      ) = parallel_self_attention(
          1,
          num_att_heads_per_partition,
          hidden_size_per_att_head,
          dropout_prob,
          batch_size,
          sequence_length,
      )
  
      (
          rank,
          hidden_size,
          tensor_model_parallel_size,
          loss,
          attention_layer,
          identity_layer,
      ) = parallel_self_attention(
          tensor_model_parallel_size,
          num_att_heads_per_partition,
          hidden_size_per_att_head,
          dropout_prob,
          batch_size,
          sequence_length,
      )
      assert hideen_size_1 == hidden_size
  
      error = loss_1.sub(loss).abs().max()
      torch.distributed.barrier()
      print(
          "   loss error on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 5.0e-6
  
      my_lin_grad_list = torch.split(
          attention_layer_1.query_key_value.weight.grad,
          hidden_size // tensor_model_parallel_size,
          0,
      )[rank::tensor_model_parallel_size]
      my_lin_grad = torch.cat(my_lin_grad_list, dim=0)
      error = (
          my_lin_grad.sub(attention_layer.query_key_value.weight.grad)
          .abs()
          .max()
      )
      torch.distributed.barrier()
      print(
          "   weight gradient error on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 5.0e-6
  
      error = (
          identity_layer_1.weight.grad.sub(identity_layer.weight.grad)
          .abs()
          .max()
      )
      torch.distributed.barrier()
      print(
          "   input gradient error on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 5.0e-6
  
      torch.distributed.barrier()
      if torch.distributed.get_rank() == 0:
          print(" >> passed the test :-)")
  
imports: from mpu import layers
  from commons import set_random_seed
  from commons import print_separator
  from commons import initialize_distributed
  import mpu
  from torch.nn.parameter import Parameter
  import torch.nn.init
  import torch
  import random
  import sys
label_desc: Use image node_id 6 for calling a global function with example usage: test_parallel_self_attention(tensor_model_parallel_size) without return types