code: def split_heads(X, nheads):
      """
      Split heads:
      1) Split (reshape) last dimension (size d_model) into nheads, d_head
      2) Transpose X from (batch size, sequence length, nheads, d_head) to
          (batch size, nheads, sequence length, d_head)
  
      Inputs:
        X : [batch size, sequence length, nheads * d_head]
        nheads : integer
      Outputs:
        [batch size,  nheads, sequence length, d_head]
  
      """
      last_dim = X.shape[-1]
      assert last_dim % nheads == 0
      X_last_dim_split = X.view(
          list(X.shape[:-1]) + [nheads, last_dim // nheads]
      )
      return X_last_dim_split.transpose(1, 2)
  
file_path: translate/pytorch_translate/research/attention/multihead_attention.py
imports: import math
  import numpy
  import torch
  import torch.nn.functional
  from torch import nn
label_desc: Use image node_id 4 for calling a global function with example usage: split_heads(X, nheads) and returns: X_last_dim_split