class_docstr: Parallel self-attention layer abstract class.
  
  Self-attention layer takes input with size [s, b, h]
  and returns output of the same size.
class_docstr_tok: ["Parallel","self-attention","layer","abstract","class",".","Self-attention","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","output","of","the","same","size","."]
code: def _transpose_last_dim(
      self, mixed_layer, num_splits, num_splits_first
  ):
      input_shape = mixed_layer.size()
      if num_splits_first:
          """[s, b, num_splits * np * hn]
          -->(view) [s, b, num_splits, np, hn]
          -->(tranpose) [s, b, np, num_splits, hn]
          -->(view) [s, b, np * num_splits * hn]"""
  
          intermediate_shape = input_shape[:-1] + (
              num_splits,
              self.num_attention_heads_per_partition,
              self.hidden_size_per_attention_head,
          )
  
          mixed_layer = mixed_layer.view(*intermediate_shape)
          mixed_layer = mixed_layer.transpose(-2, -3).contiguous()
      else:
          """[s, b, np * hn * num_splits]
          -->(view) [s, b, np, hn, num_splits]
          -->(tranpose) [s, b, np, num_splits, hn]
          -->(view) [s, b, np * num_splits * hn]"""
  
          intermediate_shape = input_shape[:-1] + (
              self.num_attention_heads_per_partition,
              self.hidden_size_per_attention_head,
              num_splits,
          )
  
          mixed_layer = mixed_layer.view(*intermediate_shape)
          mixed_layer = mixed_layer.transpose(-1, -2).contiguous()
      mixed_layer = mixed_layer.view(*input_shape)
  
      return mixed_layer
  
imports: import math
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, InfusedAdapterConfig, LoraKQVAdapterConfig, LoraKQVAdapterWeightTyingConfig, LoraKVAdapterConfig, LoraQAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_softmax import MatchedScaleMaskSoftmax
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import XPOSPositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.position_embedding.rotary_position_embedding import apply_rotary_pos_emb
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, _cast_if_autocast_enabled, attention_mask_func
  from nemo.core import adapter_mixins
label_desc: Use image node_id 4 for calling the ParallelAttention obj's underlying member method code with example usage: obj._transpose_last_dim(mixed_layer, num_splits, num_splits_first) and returns: mixed_layer