code: import tensorflow as tf
  import numpy as np
  import math
  from tacotron.utils.ops import shape_list
  
  class MultiheadAttention():
    '''Computes the multi-head attention as described in
    https://arxiv.org/abs/1706.03762.
    Args:
      num_heads: The number of attention heads.
      query: The sequence of queries. A tensor of shape :math:`[B, T_1, ...]`.
      value: The sequence to attend. A tensor of shape :math:`[B, T_2, ...]`.
        If ``None``, computes self-attention.
      num_units: The number of hidden units. If not set, it is set to the input
        dimension.
      attention_type: a string, either "dot_attention", "mlp_attention".
    Returns:
       The concatenated attention context of each head.
    '''
    def __init__(self,
                 query,
                 value,
                 num_heads=4,
                 attention_type='mlp_attention',
                 num_units=None,
                 normalize=True):
      self.query = query
      self.value = value
      self.num_heads = num_heads
      self.attention_type = attention_type
      self.num_units = num_units or query.get_shape().as_list()[-1]
      self.normalize = normalize
  
    def multi_head_attention(self):
      if self.num_units % self.num_heads != 0:
        raise ValueError("Multi head attention requires that num_units is a"
                         " multiple of {}".format(num_heads))
  
      with tf.variable_scope("Multihead-attention"):
        q = tf.layers.conv1d(self.query, self.num_units, 1)
        k =  tf.layers.conv1d(self.value, self.num_units, 1)
        v = self.value
        qs, ks, vs = self._split_heads(q, k, v)
        if self.attention_type == 'mlp_attention':
          style_embeddings = self._mlp_attention(qs, ks, vs)
        elif self.attention_type == 'dot_attention':
          style_embeddings = self._dot_product(qs, ks, vs)
        else:
          raise ValueError('Only mlp_attention and dot_attention are supported')
  
        return self._combine_heads(style_embeddings)
  
    def _split_heads(self, q, k, v):
      '''Split the channels into multiple heads
      
      Returns:
           Tensors with shape [batch, num_heads, length_x, dim_x/num_heads]
      '''
      qs = tf.transpose(self._split_last_dimension(q, self.num_heads), [0, 2, 1, 3])
      ks = tf.transpose(self._split_last_dimension(k, self.num_heads), [0, 2, 1, 3])
      v_shape = shape_list(v)  
      vs = tf.tile(tf.expand_dims(v, axis=1), [1, self.num_heads, 1, 1])
      return qs, ks, vs
  
    def _split_last_dimension(self, x, num_heads):
      '''Reshape x to num_heads
      Returns:
          a Tensor with shape [batch, length_x, num_heads, dim_x/num_heads]
      '''
      x_shape = shape_list(x)
      dim = x_shape[-1]
      assert dim % num_heads == 0 
      return tf.reshape(x, x_shape[:-1] + [num_heads, dim // num_heads])
  
    def _dot_product(self, qs, ks, vs):
      '''dot-product computation
      Returns:
          a context vector with shape [batch, num_heads, length_q, dim_vs]
      '''
      qk = tf.matmul(qs, ks, transpose_b=True)
      scale_factor = (self.num_units // self.num_heads)**-0.5
      if self.normalize:
        qk *= scale_factor
      weights = tf.nn.softmax(qk, name="dot_attention_weights")
      context = tf.matmul(weights, vs)
      return context
  
    def _mlp_attention(self, qs, ks, vs):
      '''MLP computation modified from https://github.com/npuichigo
      Returns:
          a context vector with shape [batch, num_heads, length_q, dim_vs]
      '''
      num_units = qs.get_shape()[-1].value
      dtype = qs.dtype
  
      v = tf.get_variable("attention_v", [num_units], dtype=dtype)
      if self.normalize:
        #https://github.com/tensorflow/tensorflow/blob/r1.7/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py#L470
        # Scalar used in weight normalization
        g = tf.get_variable(
            "attention_g", dtype=dtype,
            initializer=math.sqrt((1. / num_units)))
        # Bias added prior to the nonlinearity
        b = tf.get_variable(
            "attention_b", [num_units], dtype=dtype,
            initializer=tf.zeros_initializer())
        # normed_v = g * v / ||v||
        normed_v = g * v * tf.rsqrt(
                tf.reduce_sum(tf.square(v)))
        # Single layer multilayer perceptron.
        add = tf.reduce_sum(normed_v * tf.tanh(ks + qs + b), [-1], keep_dims=True)
      else:
        # Single layer multilayer perceptron.
        add = tf.reduce_sum(v * tf.tanh(ks + qs), [-1], keep_dims=True)
  
      # Compute attention weights.
      weights = tf.nn.softmax(tf.transpose(add, [0, 1, 3, 2]), name="mlp_attention_weights")
      # Compute attention context.
      context = tf.matmul(weights, vs)
      return context
  
    def _combine_heads(self, x):
      '''Combine all heads
         Returns:
             a Tensor with shape [batch, length_x, shape_x[-1] * shape_x[-3]]
      '''
      x = tf.transpose(x, [0, 2, 1, 3])
      x_shape = shape_list(x)
      return tf.reshape(x, x_shape[:-2] + [self.num_heads * x_shape[-1]])
    
apis: tensorflow.layers.conv1d tensorflow.matmul tensorflow.nn.softmax tensorflow.transpose tensorflow.get_variable tensorflow.zeros_initializer tensorflow.reshape tensorflow.expand_dims tensorflow.tanh tensorflow.square tensorflow.variable_scope
api_extract: [(61, 'tacotron.utils.ops.shape_list', 'shape_list', (['v'], {}), False, 'from tacotron.utils.ops import shape_list\n'), (70, 'tacotron.utils.ops.shape_list', 'shape_list', (['x'], {}), False, 'from tacotron.utils.ops import shape_list\n'), (73, 'tensorflow.reshape', 'tf.reshape', (['x', '(x_shape[:-1] + [num_heads, dim // num_heads])'], {}), True, 'import tensorflow as tf\n'), (80, 'tensorflow.matmul', 'tf.matmul', (['qs', 'ks'], {'transpose_b': '(True)'}), True, 'import tensorflow as tf\n'), (84, 'tensorflow.nn.softmax', 'tf.nn.softmax', (['qk'], {'name': '"""dot_attention_weights"""'}), True, 'import tensorflow as tf\n'), (85, 'tensorflow.matmul', 'tf.matmul', (['weights', 'vs'], {}), True, 'import tensorflow as tf\n'), (96, 'tensorflow.get_variable', 'tf.get_variable', (['"""attention_v"""', '[num_units]'], {'dtype': 'dtype'}), True, 'import tensorflow as tf\n'), (119, 'tensorflow.matmul', 'tf.matmul', (['weights', 'vs'], {}), True, 'import tensorflow as tf\n'), (127, 'tensorflow.transpose', 'tf.transpose', (['x', '[0, 2, 1, 3]'], {}), True, 'import tensorflow as tf\n'), (128, 'tacotron.utils.ops.shape_list', 'shape_list', (['x'], {}), False, 'from tacotron.utils.ops import shape_list\n'), (129, 'tensorflow.reshape', 'tf.reshape', (['x', '(x_shape[:-2] + [self.num_heads * x_shape[-1]])'], {}), True, 'import tensorflow as tf\n'), (39, 'tensorflow.variable_scope', 'tf.variable_scope', (['"""Multihead-attention"""'], {}), True, 'import tensorflow as tf\n'), (40, 'tensorflow.layers.conv1d', 'tf.layers.conv1d', (['self.query', 'self.num_units', '(1)'], {}), True, 'import tensorflow as tf\n'), (41, 'tensorflow.layers.conv1d', 'tf.layers.conv1d', (['self.value', 'self.num_units', '(1)'], {}), True, 'import tensorflow as tf\n'), (62, 'tensorflow.expand_dims', 'tf.expand_dims', (['v'], {'axis': '(1)'}), True, 'import tensorflow as tf\n'), (117, 'tensorflow.transpose', 'tf.transpose', (['add', '[0, 1, 3, 2]'], {}), True, 'import tensorflow as tf\n'), (102, 'math.sqrt', 'math.sqrt', (['(1.0 / num_units)'], {}), False, 'import math\n'), (106, 'tensorflow.zeros_initializer', 'tf.zeros_initializer', ([], {}), True, 'import tensorflow as tf\n'), (111, 'tensorflow.tanh', 'tf.tanh', (['(ks + qs + b)'], {}), True, 'import tensorflow as tf\n'), (114, 'tensorflow.tanh', 'tf.tanh', (['(ks + qs)'], {}), True, 'import tensorflow as tf\n'), (109, 'tensorflow.square', 'tf.square', (['v'], {}), True, 'import tensorflow as tf\n')]