class_docstr:  Multi-headed attention is based on "Attention
  is all you Need" (https://arxiv.org/pdf/1706.03762.pdf).
class_docstr_tok: ["Multi-headed","attention","is","based","on","``","Attention","is","all","you","Need","''","(","https",":","\/\/arxiv.org\/pdf\/1706.03762.pdf",")","."]
docstr: Split hidden_size into depth(hidden_size // num_heads) for
  multi-head attention.
  Args:
    x: (batch_size, seq_len_x, hidden_size)
    batch_size
  
  Returns:
    split_x: (batch_size, num_heads, seq_len_x, depth)
docstr_tok: ["Split","hidden_size","into","depth","(","hidden_size","\/\/","num_heads",")","for","multi-head","attention",".","Args",":","x",":","(","batch_size",",","seq_len_x",",","hidden_size",")","batch_size","Returns",":","split_x",":","(","batch_size",",","num_heads",",","seq_len_x",",","depth",")"]
code: def split_heads(self, x, batch_size):
      """
      Split hidden_size into depth(hidden_size // num_heads) for
      multi-head attention.
      Args:
        x: (batch_size, seq_len_x, hidden_size)
        batch_size
  
      Returns:
        split_x: (batch_size, num_heads, seq_len_x, depth)
      """
      x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
      split_x = tf.transpose(x, perm=[0, 2, 1, 3])
      return split_x
  
file_path: delta/delta/layers/sub_tf.py
imports: from absl import logging
  import delta.compat
  import numpy
  from delta.layers.base_layer import Layer
  from delta.layers.utils_tf import shape_list
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.split_heads(x, batch_size) and returns: split_x