easy_rec/python/layers/multihead_cross_attention.py

# -*- encoding:utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math

import tensorflow as tf

from easy_rec.python.compat.layers import layer_norm as tf_layer_norm
from easy_rec.python.utils.activation import gelu
from easy_rec.python.utils.shape_utils import get_shape_list

if tf.__version__ >= '2.0':
  tf = tf.compat.v1


def create_initializer(initializer_range=0.02):
  """Creates a `truncated_normal_initializer` with the given range."""
  return tf.truncated_normal_initializer(stddev=initializer_range)


def dropout(input_tensor, dropout_prob):
  """Perform dropout.

  Args:
    input_tensor: float Tensor.
    dropout_prob: Python float. The probability of dropping out a value (NOT of
      *keeping* a dimension as in `tf.nn.dropout`).

  Returns:
    A version of `input_tensor` with dropout applied.
  """
  if dropout_prob is None or dropout_prob == 0.0:
    return input_tensor

  output = tf.nn.dropout(input_tensor, 1.0 - dropout_prob)
  return output


def attention_layer(from_tensor,
                    to_tensor,
                    size_per_head,
                    num_attention_heads=1,
                    attention_mask=None,
                    query_act=None,
                    key_act=None,
                    value_act=None,
                    attention_probs_dropout_prob=0.0,
                    initializer_range=0.02,
                    do_return_2d_tensor=False,
                    batch_size=None,
                    from_seq_length=None,
                    to_seq_length=None,
                    reuse=None):
  """Performs multi-headed attention from `from_tensor` to `to_tensor`.

  This is an implementation of multi-headed attention based on "Attention is all you Need".
  If `from_tensor` and `to_tensor` are the same, then this is self-attention.
  Each timestep in `from_tensor` attends to the corresponding sequence in `to_tensor`,
  and returns a fixed-width vector.
  This function first projects `from_tensor` into a "query" tensor and `to_tensor` into "key" and "value" tensors.
  These are (effectively) a list of tensors of length `num_attention_heads`, where each tensor is of shape:
  [batch_size, seq_length, size_per_head].
  Then, the query and key tensors are dot-producted and scaled. These are
  softmaxed to obtain attention probabilities. The value tensors are then
  interpolated by these probabilities, then concatenated back to a single
  tensor and returned.
  In practice, the multi-headed attention are done with transposes and
  reshapes rather than actual separate tensors.

  Args:
    from_tensor: float Tensor of shape [batch_size, from_seq_length,
      from_width].
    to_tensor: float Tensor of shape [batch_size, to_seq_length, to_width].
    size_per_head: int. Size of each attention head.
    num_attention_heads: int. Number of attention heads.
    attention_mask: (optional) int32 Tensor of shape [batch_size,
      from_seq_length, to_seq_length]. The values should be 1 or 0. The
      attention scores will effectively be set to -infinity for any positions in
      the mask that are 0, and will be unchanged for positions that are 1.
    query_act: (optional) Activation function for the query transform.
    key_act: (optional) Activation function for the key transform.
    value_act: (optional) Activation function for the value transform.
    attention_probs_dropout_prob: (optional) float. Dropout probability of the
      attention probabilities.
    initializer_range: float. Range of the weight initializer.
    do_return_2d_tensor: bool. If True, the output will be of shape [batch_size
      * from_seq_length, num_attention_heads * size_per_head]. If False, the
      output will be of shape [batch_size, from_seq_length, num_attention_heads
      * size_per_head].
    batch_size: (Optional) int. If the input is 2D, this might be the batch size
      of the 3D version of the `from_tensor` and `to_tensor`.
    from_seq_length: (Optional) If the input is 2D, this might be the seq length
      of the 3D version of the `from_tensor`.
    to_seq_length: (Optional) If the input is 2D, this might be the seq length
      of the 3D version of the `to_tensor`.
    reuse: whether to reuse this layer

  Returns:
    float Tensor of shape [batch_size, from_seq_length,
      num_attention_heads * size_per_head]. (If `do_return_2d_tensor` is
      true, this will be of shape [batch_size * from_seq_length,
      num_attention_heads * size_per_head]).

  Raises:
    ValueError: Any of the arguments or tensor shapes are invalid.
  """

  def transpose_for_scores(input_tensor, batch_size, num_attention_heads,
                           seq_length, width):
    output_tensor = tf.reshape(
        input_tensor, [batch_size, seq_length, num_attention_heads, width])

    output_tensor = tf.transpose(output_tensor, [0, 2, 1, 3])
    return output_tensor

  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3])
  to_shape = get_shape_list(to_tensor, expected_rank=[2, 3])

  if len(from_shape) != len(to_shape):
    raise ValueError(
        'The rank of `from_tensor` must match the rank of `to_tensor`.')

  if len(from_shape) == 3:
    batch_size = from_shape[0]
    from_seq_length = from_shape[1]
    to_seq_length = to_shape[1]
  elif len(from_shape) == 2:
    if (batch_size is None or from_seq_length is None or to_seq_length is None):
      raise ValueError(
          'When passing in rank 2 tensors to attention_layer, the values '
          'for `batch_size`, `from_seq_length`, and `to_seq_length` '
          'must all be specified.')

  # Scalar dimensions referenced here:
  #   B = batch size (number of sequences)
  #   F = `from_tensor` sequence length
  #   T = `to_tensor` sequence length
  #   N = `num_attention_heads`
  #   H = `size_per_head`

  from_tensor_2d = reshape_to_matrix(from_tensor)
  to_tensor_2d = reshape_to_matrix(to_tensor)

  # `query_layer` = [B*F, N*H]
  query_layer = tf.layers.dense(
      from_tensor_2d,
      num_attention_heads * size_per_head,
      activation=query_act,
      name='query',
      kernel_initializer=create_initializer(initializer_range),
      reuse=reuse)

  # `key_layer` = [B*T, N*H]
  key_layer = tf.layers.dense(
      to_tensor_2d,
      num_attention_heads * size_per_head,
      activation=key_act,
      name='key',
      kernel_initializer=create_initializer(initializer_range),
      reuse=reuse)

  # `value_layer` = [B*T, N*H]
  value_layer = tf.layers.dense(
      to_tensor_2d,
      num_attention_heads * size_per_head,
      activation=value_act,
      name='value',
      kernel_initializer=create_initializer(initializer_range),
      reuse=reuse)

  # `query_layer` = [B, N, F, H]
  query_layer = transpose_for_scores(query_layer, batch_size,
                                     num_attention_heads, from_seq_length,
                                     size_per_head)

  # `key_layer` = [B, N, T, H]
  key_layer = transpose_for_scores(key_layer, batch_size, num_attention_heads,
                                   to_seq_length, size_per_head)

  # Take the dot product between "query" and "key" to get the raw
  # attention scores.
  # `attention_scores` = [B, N, F, T]
  attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
  attention_scores = tf.multiply(attention_scores,
                                 1.0 / math.sqrt(float(size_per_head)))

  if attention_mask is not None:
    # `attention_mask` = [B, 1, F, T]
    attention_mask = tf.expand_dims(attention_mask, axis=[1])

    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and -10000.0 for masked positions.
    adder = (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0

    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    attention_scores += adder

  # Normalize the attention scores to probabilities.
  # `attention_probs` = [B, N, F, T]
  attention_probs = tf.nn.softmax(attention_scores)

  # This is actually dropping out entire tokens to attend to, which might
  # seem a bit unusual, but is taken from the original Transformer paper.
  attention_probs = dropout(attention_probs, attention_probs_dropout_prob)

  # `value_layer` = [B, T, N, H]
  value_layer = tf.reshape(
      value_layer,
      [batch_size, to_seq_length, num_attention_heads, size_per_head])

  # `value_layer` = [B, N, T, H]
  value_layer = tf.transpose(value_layer, [0, 2, 1, 3])

  # `context_layer` = [B, N, F, H]
  context_layer = tf.matmul(attention_probs, value_layer)

  # `context_layer` = [B, F, N, H]
  context_layer = tf.transpose(context_layer, [0, 2, 1, 3])

  if do_return_2d_tensor:
    # `context_layer` = [B*F, N*H]
    context_layer = tf.reshape(
        context_layer,
        [batch_size * from_seq_length, num_attention_heads * size_per_head])
  else:
    # `context_layer` = [B, F, N*H]
    context_layer = tf.reshape(
        context_layer,
        [batch_size, from_seq_length, num_attention_heads * size_per_head])

  return context_layer


def transformer_encoder(input_tensor,
                        attention_mask=None,
                        hidden_size=768,
                        num_hidden_layers=12,
                        num_attention_heads=12,
                        intermediate_size=3072,
                        intermediate_act_fn=gelu,
                        hidden_dropout_prob=0.1,
                        attention_probs_dropout_prob=0.1,
                        initializer_range=0.02,
                        reuse=None,
                        name='transformer'):
  """Multi-headed, multi-layer Transformer from "Attention is All You Need".

  This is almost an exact implementation of the original Transformer encoder.
  See the original paper:
  https://arxiv.org/abs/1706.03762
  Args:
    input_tensor: float Tensor of shape [batch_size, seq_length, hidden_size].
    attention_mask: (optional) int32 Tensor of shape [batch_size, seq_length,
      seq_length], with 1 for positions that can be attended to and 0 in
      positions that should not be.
    hidden_size: int. Hidden size of the Transformer.
    num_hidden_layers: int. Number of layers (blocks) in the Transformer.
    num_attention_heads: int. Number of attention heads in the Transformer.
    intermediate_size: int. The size of the "intermediate" (a.k.a., feed
      forward) layer.
    intermediate_act_fn: function. The non-linear activation function to apply
      to the output of the intermediate/feed-forward layer.
    hidden_dropout_prob: float. Dropout probability for the hidden layers.
    attention_probs_dropout_prob: float. Dropout probability of the attention
      probabilities.
    initializer_range: float. Range of the initializer (stddev of truncated
      normal).
    reuse: whether to reuse this encoder
    name: scope name prefix

  Returns:
    float Tensor of shape [batch_size, seq_length, hidden_size], the final
    hidden layer of the Transformer.

  Raises:
    ValueError: A Tensor shape or parameter is invalid.
  """
  if hidden_size % num_attention_heads != 0:
    raise ValueError(
        'The hidden size (%d) is not a multiple of the number of attention '
        'heads (%d)' % (hidden_size, num_attention_heads))

  attention_head_size = int(hidden_size / num_attention_heads)
  input_shape = get_shape_list(input_tensor, expected_rank=3)
  batch_size = input_shape[0]
  seq_length = input_shape[1]
  input_width = input_shape[2]

  # The Transformer performs sum residuals on all layers so the input needs
  # to be the same as the hidden size.
  if input_width != hidden_size:
    raise ValueError('The width of the input tensor (%d) != hidden size (%d)' %
                     (input_width, hidden_size))

  # We keep the representation as a 2D tensor to avoid re-shaping it back and
  # forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on
  # the GPU/CPU but may not be free on the TPU, so we want to minimize them to
  # help the optimizer.
  prev_output = reshape_to_matrix(input_tensor)

  for layer_idx in range(num_hidden_layers):
    with tf.variable_scope('%s_layer_%d' % (name, layer_idx)):
      layer_input = prev_output

      with tf.variable_scope('attention'):
        with tf.variable_scope('self'):
          # [batch_size * from_seq_length, num_attention_heads * size_per_head]
          attention_output = attention_layer(
              from_tensor=layer_input,
              to_tensor=layer_input,
              size_per_head=attention_head_size,
              num_attention_heads=num_attention_heads,
              attention_mask=attention_mask,
              attention_probs_dropout_prob=attention_probs_dropout_prob,
              initializer_range=initializer_range,
              do_return_2d_tensor=True,
              batch_size=batch_size,
              from_seq_length=seq_length,
              to_seq_length=seq_length,
              reuse=reuse)

        # Run a linear projection of `hidden_size` then add a residual
        # with `layer_input`.
        with tf.variable_scope('output', reuse=reuse):
          attention_output = tf.layers.dense(
              attention_output,
              hidden_size,
              kernel_initializer=create_initializer(initializer_range))
          attention_output = dropout(attention_output, hidden_dropout_prob)
          attention_output = layer_norm(attention_output + layer_input)

      # The activation is only applied to the "intermediate" hidden layer.
      with tf.variable_scope('intermediate', reuse=reuse):
        intermediate_output = tf.layers.dense(
            attention_output,
            intermediate_size,
            activation=intermediate_act_fn,
            kernel_initializer=create_initializer(initializer_range))

      # Down-project back to `hidden_size` then add the residual.
      with tf.variable_scope('output', reuse=reuse):
        layer_output = tf.layers.dense(
            intermediate_output,
            hidden_size,
            kernel_initializer=create_initializer(initializer_range))
        layer_output = dropout(layer_output, hidden_dropout_prob)
        layer_output = layer_norm(layer_output + attention_output)
        prev_output = layer_output

  final_output = reshape_from_matrix(prev_output, input_shape)
  return final_output


def cross_attention_block(from_tensor,
                          to_tensor,
                          layer_idx,
                          size_per_head,
                          cross_attention_mask=None,
                          self_attention_mask=None,
                          num_attention_heads=1,
                          intermediate_size=512,
                          hidden_dropout_prob=0.1,
                          attention_probs_dropout_prob=0.1,
                          initializer_range=0.02,
                          name=''):
  """Multi-headed cross attention block.

    Args:
    from_tensor: float Tensor of shape [batch_size, from_seq_length,
      from_width].
    to_tensor: float Tensor of shape [batch_size, to_seq_length, to_width].
    layer_idx: int. layer id in the Transformer.
    size_per_head: int. Size of each attention head.
    cross_attention_mask: (optional) int32 Tensor of shape [batch_size, from_seq_length,
      to_seq_length], with 1 for positions that can be attended to and 0 in
      positions that should not be.
    self_attention_mask: (optional) int32 Tensor of shape [batch_size, from_seq_length,
      from_seq_length], with 1 for positions that can be attended to and 0 in
      positions that should not be.
    num_attention_heads: int. Number of attention heads in the Transformer.
    intermediate_size: int. The size of the "intermediate" (a.k.a., feed
      forward) layer.
    hidden_dropout_prob: float. Dropout probability for the hidden layers.
    attention_probs_dropout_prob: float. Dropout probability of the attention
      probabilities.
    initializer_range: float. Range of the initializer (stddev of truncated
      normal).
    name: scope name prefix

  Returns:
    float Tensor of shape [batch_size, seq_length, hidden_size], the final
    hidden layer of the Transformer.

  Raises:
    ValueError: A Tensor shape or parameter is invalid.
  """
  input_shape = get_shape_list(from_tensor, expected_rank=3)
  batch_size = input_shape[0]
  from_seq_length = input_shape[1]

  input_shape = get_shape_list(to_tensor, expected_rank=3)
  to_seq_length = input_shape[1]

  with tf.variable_scope('%scross_layer_%d' % (name, layer_idx)):
    with tf.variable_scope('attention'):
      with tf.variable_scope('cross'):
        # [batch_size * from_seq_length, num_attention_heads * size_per_head]
        cross_attention_output = attention_layer(
            from_tensor=from_tensor,
            to_tensor=to_tensor,
            size_per_head=size_per_head,
            num_attention_heads=num_attention_heads,
            attention_mask=cross_attention_mask,
            attention_probs_dropout_prob=attention_probs_dropout_prob,
            initializer_range=initializer_range,
            do_return_2d_tensor=True,
            batch_size=batch_size,
            from_seq_length=from_seq_length,
            to_seq_length=to_seq_length)

      with tf.variable_scope('self'):
        # [batch_size * from_seq_length, num_attention_heads * size_per_head]
        self_attention_output = attention_layer(
            from_tensor=cross_attention_output,
            to_tensor=cross_attention_output,
            size_per_head=size_per_head,
            num_attention_heads=num_attention_heads,
            attention_mask=self_attention_mask,
            attention_probs_dropout_prob=attention_probs_dropout_prob,
            initializer_range=initializer_range,
            do_return_2d_tensor=True,
            batch_size=batch_size,
            from_seq_length=from_seq_length,
            to_seq_length=from_seq_length)

      with tf.variable_scope('output'):
        attention_output = dropout(self_attention_output, hidden_dropout_prob)
        attention_output = layer_norm(attention_output + cross_attention_output)

    # The activation is only applied to the "intermediate" hidden layer.
    with tf.variable_scope('intermediate'):
      intermediate_output = tf.layers.dense(
          attention_output,
          intermediate_size,
          activation=tf.nn.relu,
          kernel_initializer=create_initializer(initializer_range))

    # Down-project back to `hidden_size` then add the residual.
    with tf.variable_scope('output'):
      layer_output = tf.layers.dense(
          intermediate_output,
          num_attention_heads * size_per_head,
          kernel_initializer=create_initializer(initializer_range))
      layer_output = dropout(layer_output, hidden_dropout_prob)
      # [batch_size * from_seq_length, num_attention_heads * size_per_head]
      layer_output = layer_norm(layer_output + attention_output)

  final_output = reshape_from_matrix(
      layer_output,
      [batch_size, from_seq_length, num_attention_heads * size_per_head])
  return final_output  # [batch_size, from_seq_length, num_attention_heads * size_per_head]


def cross_attention_tower(left_tensor,
                          right_tensor,
                          num_hidden_layers=1,
                          num_attention_heads=12,
                          left_size_per_head=64,
                          right_size_per_head=64,
                          left_intermediate_size=0,
                          right_intermediate_size=0,
                          left_input_mask=None,
                          right_input_mask=None,
                          hidden_dropout_prob=0.1,
                          attention_probs_dropout_prob=0.1,
                          initializer_range=0.02,
                          name=''):
  """Multi-headed, multi layer cross attention block.

    Args:
    left_tensor: float Tensor of shape [batch_size, left_seq_length,
      from_width].
    right_tensor: float Tensor of shape [batch_size, right_seq_length, to_width].
    num_hidden_layers: int. Number of layers (blocks) in the Transformer.
    num_attention_heads: int. Number of attention heads in the Transformer.
    left_size_per_head: int. Size of each attention head of left tower.
    right_size_per_head: int. Size of each attention head of right tower.
    left intermediate_size: int. The size of the "intermediate" (a.k.a., feed
      forward) layer of left tower. Less or equal to 0 means `num_attention_heads
      * left_size_per_head`
    right intermediate_size: int. The size of the "intermediate" (a.k.a., feed
      forward) layer of right tower. Less or equal to 0 means `num_attention_heads
      * right_size_per_head`
    left_input_mask: the mask for `left_tensor`
    right_input_mask: the mask for `right_tensor`
    hidden_dropout_prob: float. Dropout probability for the hidden layers.
    attention_probs_dropout_prob: float. Dropout probability of the attention
      probabilities.
    initializer_range: float. Range of the initializer (stddev of truncated
      normal).
    name: scope name prefix

  Returns:
    tuple of float Tensors of shape ([batch_size, left_seq_length, hidden_size],
      [batch_size, right_seq_length, hidden_size]),
      where hidden_size = num_attention_heads * size_per_head

  Raises:
    ValueError: A Tensor shape or parameter is invalid.
  """
  if left_intermediate_size <= 0:
    left_intermediate_size = num_attention_heads * left_size_per_head
  if right_intermediate_size <= 0:
    right_intermediate_size = num_attention_heads * right_size_per_head

  left_attention_mask = None
  if left_input_mask is not None:
    left_attention_mask = create_attention_mask_from_input_mask(
        left_tensor, left_attention_mask)

  left_2_right_attention_mask = None
  if right_input_mask is not None:
    left_2_right_attention_mask = create_attention_mask_from_input_mask(
        left_tensor, right_input_mask)

  right_attention_mask = None
  if right_input_mask is not None:
    right_attention_mask = create_attention_mask_from_input_mask(
        right_tensor, right_input_mask)

  right_2_left_attention_mask = None
  if left_input_mask is not None:
    right_2_left_attention_mask = create_attention_mask_from_input_mask(
        right_tensor, left_input_mask)

  prev_left_output = left_tensor
  prev_right_output = right_tensor
  for layer_idx in range(num_hidden_layers):
    left_output = cross_attention_block(
        prev_left_output,
        prev_right_output,
        layer_idx,
        num_attention_heads=num_attention_heads,
        size_per_head=left_size_per_head,
        intermediate_size=left_intermediate_size,
        hidden_dropout_prob=hidden_dropout_prob,
        cross_attention_mask=left_2_right_attention_mask,
        self_attention_mask=left_attention_mask,
        attention_probs_dropout_prob=attention_probs_dropout_prob,
        initializer_range=initializer_range,
        name='%sleft_to_right_' % name)
    right_output = cross_attention_block(
        prev_right_output,
        prev_left_output,
        layer_idx,
        num_attention_heads=num_attention_heads,
        size_per_head=right_size_per_head,
        intermediate_size=right_intermediate_size,
        hidden_dropout_prob=hidden_dropout_prob,
        cross_attention_mask=right_2_left_attention_mask,
        self_attention_mask=right_attention_mask,
        attention_probs_dropout_prob=attention_probs_dropout_prob,
        initializer_range=initializer_range,
        name='%sright_to_left_' % name)
    prev_left_output = left_output
    prev_right_output = right_output
  return prev_left_output, prev_right_output


def layer_norm(input_tensor, name=None):
  """Run layer normalization on the last dimension of the tensor."""
  return tf_layer_norm(
      inputs=input_tensor, begin_norm_axis=-1, begin_params_axis=-1, scope=name)


def reshape_to_matrix(input_tensor):
  """Reshapes a >= rank 2 tensor to a rank 2 tensor (i.e., a matrix)."""
  ndims = input_tensor.shape.ndims
  if ndims < 2:
    raise ValueError('Input tensor must have at least rank 2. Shape = %s' %
                     (input_tensor.shape))
  if ndims == 2:
    return input_tensor

  width = input_tensor.shape[-1]
  output_tensor = tf.reshape(input_tensor, [-1, width])
  return output_tensor


def reshape_from_matrix(output_tensor, orig_shape_list):
  """Reshapes a rank 2 tensor back to its original rank >= 2 tensor."""
  if len(orig_shape_list) == 2:
    return output_tensor

  output_shape = get_shape_list(output_tensor)

  orig_dims = orig_shape_list[0:-1]
  width = output_shape[-1]

  return tf.reshape(output_tensor, orig_dims + [width])


def create_attention_mask_from_input_mask(from_tensor, to_mask):
  """Create 3D attention mask from a 2D tensor mask.

  Args:
    from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...].
    to_mask: int32 Tensor of shape [batch_size, to_seq_length].

  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3])
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]

  to_shape = get_shape_list(to_mask, expected_rank=2)
  to_seq_length = to_shape[1]

  to_mask = tf.cast(
      tf.reshape(to_mask, [batch_size, 1, to_seq_length]), tf.float32)

  # We don't assume that `from_tensor` is a mask (although it could be). We
  # don't actually care if we attend *from* padding tokens (only *to* padding)
  # tokens so we create a tensor of all ones.
  #
  # `broadcast_ones` = [batch_size, from_seq_length, 1]
  broadcast_ones = tf.ones(
      shape=tf.stack([batch_size, from_seq_length, 1]), dtype=tf.float32)

  # Here we broadcast along two dimensions to create the mask.
  mask = broadcast_ones * to_mask

  return mask


def embedding_postprocessor(input_tensor,
                            use_token_type=False,
                            token_type_ids=None,
                            token_type_vocab_size=16,
                            token_type_embedding_name='token_type_embeddings',
                            reuse_token_type=None,
                            use_position_embeddings=True,
                            position_embedding_name='position_embeddings',
                            reuse_position_embedding=None,
                            initializer_range=0.02,
                            max_position_embeddings=512,
                            dropout_prob=0.1):
  """Performs various post-processing on a word embedding tensor.

  Args:
    input_tensor: float Tensor of shape [batch_size, seq_length,
      embedding_size].
    use_token_type: bool. Whether to add embeddings for `token_type_ids`.
    token_type_ids: (optional) int32 Tensor of shape [batch_size, seq_length].
      Must be specified if `use_token_type` is True.
    token_type_vocab_size: int. The vocabulary size of `token_type_ids`.
    token_type_embedding_name: string. The name of the embedding table variable
      for token type ids.
    reuse_token_type: bool. Whether to reuse token type embedding variable.
    use_position_embeddings: bool. Whether to add position embeddings for the
      position of each token in the sequence.
    position_embedding_name: string. The name of the embedding table variable
      for positional embeddings.
    reuse_position_embedding: bool. Whether to reuse position embedding variable.
    initializer_range: float. Range of the weight initialization.
    max_position_embeddings: int. Maximum sequence length that might ever be
      used with this model. This can be longer than the sequence length of
      input_tensor, but cannot be shorter.
    dropout_prob: float. Dropout probability applied to the final output tensor.

  Returns:
    float tensor with same shape as `input_tensor`.

  Raises:
    ValueError: One of the tensor shapes or input values is invalid.
  """
  input_shape = get_shape_list(input_tensor, expected_rank=3)
  batch_size = input_shape[0]
  seq_length = input_shape[1]
  width = input_shape[2]

  output = input_tensor

  if use_token_type:
    if token_type_ids is None:
      raise ValueError('`token_type_ids` must be specified if'
                       '`use_token_type` is True.')
    with tf.variable_scope('token_type', reuse=reuse_token_type):
      token_type_table = tf.get_variable(
          name=token_type_embedding_name,
          shape=[token_type_vocab_size, width],
          initializer=create_initializer(initializer_range))
    # This vocab will be small so we always do one-hot here, since it is always
    # faster for a small vocabulary.
    flat_token_type_ids = tf.reshape(token_type_ids, [-1])
    one_hot_ids = tf.one_hot(flat_token_type_ids, depth=token_type_vocab_size)
    token_type_embeddings = tf.matmul(one_hot_ids, token_type_table)
    token_type_embeddings = tf.reshape(token_type_embeddings,
                                       [batch_size, seq_length, width])
    output += token_type_embeddings

  if use_position_embeddings:
    assert_op = tf.assert_less_equal(seq_length, max_position_embeddings)
    with tf.control_dependencies([assert_op]):
      with tf.variable_scope(
          'position_embedding', reuse=reuse_position_embedding):
        full_position_embeddings = tf.get_variable(
            name=position_embedding_name,
            shape=[max_position_embeddings, width],
            initializer=create_initializer(initializer_range))
      # Since the position embedding table is a learned variable, we create it
      # using a (long) sequence length `max_position_embeddings`. The actual
      # sequence length might be shorter than this, for faster training of
      # tasks that do not have long sequences.
      #
      # So `full_position_embeddings` is effectively an embedding table
      # for position [0, 1, 2, ..., max_position_embeddings-1], and the current
      # sequence has positions [0, 1, 2, ... seq_length-1], so we can just
      # perform a slice.
      position_embeddings = tf.slice(full_position_embeddings, [0, 0],
                                     [seq_length, -1])
      num_dims = len(output.shape.as_list())

      # Only the last two dimensions are relevant (`seq_length` and `width`), so
      # we broadcast among the first dimensions, which is typically just
      # the batch size.
      position_broadcast_shape = []
      for _ in range(num_dims - 2):
        position_broadcast_shape.append(1)
      position_broadcast_shape.extend([seq_length, width])
      position_embeddings = tf.reshape(position_embeddings,
                                       position_broadcast_shape)
      output += position_embeddings

  output = layer_norm_and_dropout(output, dropout_prob)
  return output


def layer_norm_and_dropout(input_tensor, dropout_prob, name=None):
  """Runs layer normalization followed by dropout."""
  output_tensor = layer_norm(input_tensor, name)
  output_tensor = dropout(output_tensor, dropout_prob)
  return output_tensor