modeling.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf


def gelu(x):
  """Gaussian Error Linear Unit.

  This is a smoother version of the RELU.
  Original paper: https://arxiv.org/abs/1606.08415
  Args:
    x: float Tensor to perform activation.

  Returns:
    `x` with the GELU activation applied.
  """
  cdf = 0.5 * (1.0 + tf.tanh(
      (np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
  return x * cdf


def embedding_lookup(x, n_token, d_embed, initializer, use_tpu=True,
                     scope='embedding', reuse=None, dtype=tf.float32):
  """TPU and GPU embedding_lookup function."""
  with tf.variable_scope(scope, reuse=reuse):
    lookup_table = tf.get_variable('lookup_table', [n_token, d_embed],
                                   dtype=dtype, initializer=initializer)
    if use_tpu:
      one_hot_idx = tf.one_hot(x, n_token, dtype=dtype)
      if one_hot_idx.shape.ndims == 2:
        return tf.einsum('in,nd->id', one_hot_idx, lookup_table), lookup_table
      else:
        return tf.einsum('ibn,nd->ibd', one_hot_idx, lookup_table), lookup_table
    else:
      return tf.nn.embedding_lookup(lookup_table, x), lookup_table


def positional_embedding(pos_seq, inv_freq, bsz=None):
  sinusoid_inp = tf.einsum('i,d->id', pos_seq, inv_freq)
  pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)
  pos_emb = pos_emb[:, None, :]

  if bsz is not None:
    pos_emb = tf.tile(pos_emb, [1, bsz, 1])

  return pos_emb


def positionwise_ffn(inp, d_model, d_inner, dropout, kernel_initializer,
                     activation_type='relu', scope='ff', is_training=True,
                     reuse=None):
  """Position-wise Feed-forward Network."""
  if activation_type == 'relu':
    activation = tf.nn.relu
  elif activation_type == 'gelu':
    activation = gelu
  else:
    raise ValueError('Unsupported activation type {}'.format(activation_type))

  output = inp
  with tf.variable_scope(scope, reuse=reuse):
    output = tf.layers.dense(output, d_inner, activation=activation,
                             kernel_initializer=kernel_initializer,
                             name='layer_1')
    output = tf.layers.dropout(output, dropout, training=is_training,
                               name='drop_1')
    output = tf.layers.dense(output, d_model,
                             kernel_initializer=kernel_initializer,
                             name='layer_2')
    output = tf.layers.dropout(output, dropout, training=is_training,
                               name='drop_2')
    output = tf.contrib.layers.layer_norm(output + inp, begin_norm_axis=-1,
                                          scope='LayerNorm')
  return output


def head_projection(h, d_model, n_head, d_head, kernel_initializer, name):
  """Project hidden states to a specific head with a 4D-shape."""
  proj_weight = tf.get_variable('{}/kernel'.format(name),
                                [d_model, n_head, d_head], dtype=h.dtype,
                                initializer=kernel_initializer)
  head = tf.einsum('ibh,hnd->ibnd', h, proj_weight)

  return head


def post_attention(h, attn_vec, d_model, n_head, d_head, dropout, is_training,
                   kernel_initializer, residual=True):
  """Post-attention processing."""
  # post-attention projection (back to `d_model`)
  proj_o = tf.get_variable('o/kernel', [d_model, n_head, d_head],
                           dtype=h.dtype, initializer=kernel_initializer)
  attn_out = tf.einsum('ibnd,hnd->ibh', attn_vec, proj_o)

  attn_out = tf.layers.dropout(attn_out, dropout, training=is_training)
  if residual:
    output = tf.contrib.layers.layer_norm(attn_out + h, begin_norm_axis=-1,
                                          scope='LayerNorm')
  else:
    output = tf.contrib.layers.layer_norm(attn_out, begin_norm_axis=-1,
                                          scope='LayerNorm')

  return output


def abs_attn_core(q_head, k_head, v_head, attn_mask, dropatt, is_training,
                  scale):
  """Core absolute positional attention operations."""

  attn_score = tf.einsum('ibnd,jbnd->ijbn', q_head, k_head)
  attn_score *= scale
  if attn_mask is not None:
    attn_score = attn_score - 1e30 * attn_mask

  # attention probability
  attn_prob = tf.nn.softmax(attn_score, 1)
  attn_prob = tf.layers.dropout(attn_prob, dropatt, training=is_training)

  # attention output
  attn_vec = tf.einsum('ijbn,jbnd->ibnd', attn_prob, v_head)

  return attn_vec


def rel_attn_core(q_head, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat,
                  r_w_bias, r_r_bias, r_s_bias, attn_mask, dropatt, is_training,
                  scale):
  """Core relative positional attention operations."""

  # content based attention score
  ac = tf.einsum('ibnd,jbnd->ijbn', q_head + r_w_bias, k_head_h)

  # position based attention score
  bd = tf.einsum('ibnd,jbnd->ijbn', q_head + r_r_bias, k_head_r)
  bd = rel_shift(bd, klen=tf.shape(ac)[1])

  # segment based attention score
  if seg_mat is None:
    ef = 0
  else:
    ef = tf.einsum('ibnd,snd->ibns', q_head + r_s_bias, seg_embed)
    ef = tf.einsum('ijbs,ibns->ijbn', seg_mat, ef)

  # merge attention scores and perform masking
  attn_score = (ac + bd + ef) * scale
  if attn_mask is not None:
    # attn_score = attn_score * (1 - attn_mask) - 1e30 * attn_mask
    attn_score = attn_score - 1e30 * attn_mask

  # attention probability
  attn_prob = tf.nn.softmax(attn_score, 1)
  attn_prob = tf.layers.dropout(attn_prob, dropatt, training=is_training)

  # attention output
  attn_vec = tf.einsum('ijbn,jbnd->ibnd', attn_prob, v_head_h)

  return attn_vec


def rel_shift(x, klen=-1):
  """perform relative shift to form the relative attention score."""
  x_size = tf.shape(x)

  x = tf.reshape(x, [x_size[1], x_size[0], x_size[2], x_size[3]])
  x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
  x = tf.reshape(x, [x_size[0], x_size[1] - 1, x_size[2], x_size[3]])
  x = tf.slice(x, [0, 0, 0, 0], [-1, klen, -1, -1])

  return x


def _create_mask(qlen, mlen, dtype=tf.float32, same_length=False):
  """create causal attention mask."""
  attn_mask = tf.ones([qlen, qlen], dtype=dtype)
  mask_u = tf.matrix_band_part(attn_mask, 0, -1)
  mask_dia = tf.matrix_band_part(attn_mask, 0, 0)
  attn_mask_pad = tf.zeros([qlen, mlen], dtype=dtype)
  ret = tf.concat([attn_mask_pad, mask_u - mask_dia], 1)
  if same_length:
    mask_l = tf.matrix_band_part(attn_mask, -1, 0)
    ret = tf.concat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], 1)

  return ret


def _cache_mem(curr_out, prev_mem, mem_len, reuse_len=None):
  """cache hidden states into memory."""
  if mem_len is None or mem_len == 0:
    return None
  else:
    if reuse_len is not None and reuse_len > 0:
      curr_out = curr_out[:reuse_len]

    if prev_mem is None:
      new_mem = curr_out[-mem_len:]
    else:
      new_mem = tf.concat([prev_mem, curr_out], 0)[-mem_len:]

  return tf.stop_gradient(new_mem)


def relative_positional_encoding(qlen, klen, d_model, clamp_len, attn_type,
                                 bi_data, bsz=None, dtype=None):
  """create relative positional encoding."""
  freq_seq = tf.range(0, d_model, 2.0)
  if dtype is not None and dtype != tf.float32:
    freq_seq = tf.cast(freq_seq, dtype=dtype)
  inv_freq = 1 / (10000 ** (freq_seq / d_model))

  if attn_type == 'bi':
    # beg, end = klen - 1, -qlen
    beg, end = klen, -qlen
  elif attn_type == 'uni':
    # beg, end = klen - 1, -1
    beg, end = klen, -1
  else:
    raise ValueError('Unknown `attn_type` {}.'.format(attn_type))

  if bi_data:
    fwd_pos_seq = tf.range(beg, end, -1.0)
    bwd_pos_seq = tf.range(-beg, -end, 1.0)

    if dtype is not None and dtype != tf.float32:
      fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
      bwd_pos_seq = tf.cast(bwd_pos_seq, dtype=dtype)

    if clamp_len > 0:
      fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -clamp_len, clamp_len)
      bwd_pos_seq = tf.clip_by_value(bwd_pos_seq, -clamp_len, clamp_len)

    if bsz is not None:
      # With bi_data, the batch size should be divisible by 2.
      assert bsz%2 == 0
      fwd_pos_emb = positional_embedding(fwd_pos_seq, inv_freq, bsz//2)
      bwd_pos_emb = positional_embedding(bwd_pos_seq, inv_freq, bsz//2)
    else:
      fwd_pos_emb = positional_embedding(fwd_pos_seq, inv_freq)
      bwd_pos_emb = positional_embedding(bwd_pos_seq, inv_freq)

    pos_emb = tf.concat([fwd_pos_emb, bwd_pos_emb], axis=1)
  else:
    fwd_pos_seq = tf.range(beg, end, -1.0)
    if dtype is not None and dtype != tf.float32:
      fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
    if clamp_len > 0:
      fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -clamp_len, clamp_len)
    pos_emb = positional_embedding(fwd_pos_seq, inv_freq, bsz)

  return pos_emb


def multihead_attn(q, k, v, attn_mask, d_model, n_head, d_head, dropout,
                   dropatt, is_training, kernel_initializer, residual=True,
                   scope='abs_attn', reuse=None):
  """Standard multi-head attention with absolute positional embedding."""

  scale = 1 / (d_head ** 0.5)
  with tf.variable_scope(scope, reuse=reuse):
    # attention heads
    q_head = head_projection(
        q, d_model, n_head, d_head, kernel_initializer, 'q')
    k_head = head_projection(
        k, d_model, n_head, d_head, kernel_initializer, 'k')
    v_head = head_projection(
        v, d_model, n_head, d_head, kernel_initializer, 'v')

    # attention vector
    attn_vec = abs_attn_core(q_head, k_head, v_head, attn_mask, dropatt,
                             is_training, scale)

    # post processing
    output = post_attention(v, attn_vec, d_model, n_head, d_head, dropout,
                            is_training, kernel_initializer, residual)

  return output



def rel_multihead_attn(h, r, r_w_bias, r_r_bias, seg_mat, r_s_bias, seg_embed,
                       attn_mask, mems, d_model, n_head, d_head, dropout,
                       dropatt, is_training, kernel_initializer,
                       scope='rel_attn', reuse=None):
  """Multi-head attention with relative positional encoding."""

  scale = 1 / (d_head ** 0.5)
  with tf.variable_scope(scope, reuse=reuse):
    if mems is not None and mems.shape.ndims > 1:
      cat = tf.concat([mems, h], 0)
    else:
      cat = h

    # content heads
    q_head_h = head_projection(
        h, d_model, n_head, d_head, kernel_initializer, 'q')
    k_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'k')
    v_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'v')

    # positional heads
    k_head_r = head_projection(
        r, d_model, n_head, d_head, kernel_initializer, 'r')

    # core attention ops
    attn_vec = rel_attn_core(
        q_head_h, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
        r_r_bias, r_s_bias, attn_mask, dropatt, is_training, scale)

    # post processing
    output = post_attention(h, attn_vec, d_model, n_head, d_head, dropout,
                            is_training, kernel_initializer)

  return output


def two_stream_rel_attn(h, g, r, mems, r_w_bias, r_r_bias, seg_mat, r_s_bias,
                        seg_embed, attn_mask_h, attn_mask_g, target_mapping,
                        d_model, n_head, d_head, dropout, dropatt, is_training,
                        kernel_initializer, scope='rel_attn'):
  """Two-stream attention with relative positional encoding."""

  scale = 1 / (d_head ** 0.5)
  with tf.variable_scope(scope, reuse=False):

    # content based attention score
    if mems is not None and mems.shape.ndims > 1:
      cat = tf.concat([mems, h], 0)
    else:
      cat = h

    # content-based key head
    k_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'k')

    # content-based value head
    v_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'v')

    # position-based key head
    k_head_r = head_projection(
        r, d_model, n_head, d_head, kernel_initializer, 'r')

    ##### h-stream
    # content-stream query head
    q_head_h = head_projection(
        h, d_model, n_head, d_head, kernel_initializer, 'q')

    # core attention ops
    attn_vec_h = rel_attn_core(
        q_head_h, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
        r_r_bias, r_s_bias, attn_mask_h, dropatt, is_training, scale)

    # post processing
    output_h = post_attention(h, attn_vec_h, d_model, n_head, d_head, dropout,
                              is_training, kernel_initializer)

  with tf.variable_scope(scope, reuse=True):
    ##### g-stream
    # query-stream query head
    q_head_g = head_projection(
        g, d_model, n_head, d_head, kernel_initializer, 'q')

    # core attention ops
    if target_mapping is not None:
      q_head_g = tf.einsum('mbnd,mlb->lbnd', q_head_g, target_mapping)
      attn_vec_g = rel_attn_core(
          q_head_g, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
          r_r_bias, r_s_bias, attn_mask_g, dropatt, is_training, scale)
      attn_vec_g = tf.einsum('lbnd,mlb->mbnd', attn_vec_g, target_mapping)
    else:
      attn_vec_g = rel_attn_core(
          q_head_g, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
          r_r_bias, r_s_bias, attn_mask_g, dropatt, is_training, scale)

    # post processing
    output_g = post_attention(g, attn_vec_g, d_model, n_head, d_head, dropout,
                              is_training, kernel_initializer)

    return output_h, output_g


def transformer_xl(inp_k, n_token, n_layer, d_model, n_head,
                d_head, d_inner, dropout, dropatt, attn_type,
                bi_data, initializer, is_training, mem_len=None,
                inp_q=None, mems=None,
                same_length=False, clamp_len=-1, untie_r=False,
                use_tpu=True, input_mask=None,
                perm_mask=None, seg_id=None, reuse_len=None,
                ff_activation='relu', target_mapping=None,
                use_bfloat16=False, scope='transformer', **kwargs):
  """
    Defines a Transformer-XL computation graph with additional
    support for XLNet.

    Args:

    inp_k: int32 Tensor in shape [len, bsz], the input token IDs.
    seg_id: int32 Tensor in shape [len, bsz], the input segment IDs.
    input_mask: float32 Tensor in shape [len, bsz], the input mask.
      0 for real tokens and 1 for padding.
    mems: a list of float32 Tensors in shape [mem_len, bsz, d_model], memory
      from previous batches. The length of the list equals n_layer.
      If None, no memory is used.
    perm_mask: float32 Tensor in shape [len, len, bsz].
      If perm_mask[i, j, k] = 0, i attend to j in batch k;
      if perm_mask[i, j, k] = 1, i does not attend to j in batch k.
      If None, each position attends to all the others.
    target_mapping: float32 Tensor in shape [num_predict, len, bsz].
      If target_mapping[i, j, k] = 1, the i-th predict in batch k is
      on the j-th token.
      Only used during pretraining for partial prediction.
      Set to None during finetuning.
    inp_q: float32 Tensor in shape [len, bsz].
      1 for tokens with losses and 0 for tokens without losses.
      Only used during pretraining for two-stream attention.
      Set to None during finetuning.

    n_layer: int, the number of layers.
    d_model: int, the hidden size.
    n_head: int, the number of attention heads.
    d_head: int, the dimension size of each attention head.
    d_inner: int, the hidden size in feed-forward layers.
    ff_activation: str, "relu" or "gelu".
    untie_r: bool, whether to untie the biases in attention.
    n_token: int, the vocab size.

    is_training: bool, whether in training mode.
    use_tpu: bool, whether TPUs are used.
    use_bfloat16: bool, use bfloat16 instead of float32.
    dropout: float, dropout rate.
    dropatt: float, dropout rate on attention probabilities.
    init: str, the initialization scheme, either "normal" or "uniform".
    init_range: float, initialize the parameters with a uniform distribution
      in [-init_range, init_range]. Only effective when init="uniform".
    init_std: float, initialize the parameters with a normal distribution
      with mean 0 and stddev init_std. Only effective when init="normal".
    mem_len: int, the number of tokens to cache.
    reuse_len: int, the number of tokens in the currect batch to be cached
      and reused in the future.
    bi_data: bool, whether to use bidirectional input pipeline.
      Usually set to True during pretraining and False during finetuning.
    clamp_len: int, clamp all relative distances larger than clamp_len.
      -1 means no clamping.
    same_length: bool, whether to use the same attention length for each token.
    summary_type: str, "last", "first", "mean", or "attn". The method
      to pool the input to get a vector representation.
    initializer: A tf initializer.
    scope: scope name for the computation graph.
  """
  tf.logging.info('memory input {}'.format(mems))
  tf_float = tf.bfloat16 if use_bfloat16 else tf.float32
  tf.logging.info('Use float type {}'.format(tf_float))

  new_mems = []
  with tf.variable_scope(scope):
    if untie_r:
      r_w_bias = tf.get_variable('r_w_bias', [n_layer, n_head, d_head],
                                 dtype=tf_float, initializer=initializer)
      r_r_bias = tf.get_variable('r_r_bias', [n_layer, n_head, d_head],
                                 dtype=tf_float, initializer=initializer)
    else:
      r_w_bias = tf.get_variable('r_w_bias', [n_head, d_head],
                                 dtype=tf_float, initializer=initializer)
      r_r_bias = tf.get_variable('r_r_bias', [n_head, d_head],
                                 dtype=tf_float, initializer=initializer)

    bsz = tf.shape(inp_k)[1]
    qlen = tf.shape(inp_k)[0]
    mlen = tf.shape(mems[0])[0] if mems is not None else 0
    klen = mlen + qlen

    ##### Attention mask
    # causal attention mask
    if attn_type == 'uni':
      attn_mask = _create_mask(qlen, mlen, tf_float, same_length)
      attn_mask = attn_mask[:, :, None, None]
    elif attn_type == 'bi':
      attn_mask = None
    else:
      raise ValueError('Unsupported attention type: {}'.format(attn_type))

    # data mask: input mask & perm mask
    if input_mask is not None and perm_mask is not None:
      data_mask = input_mask[None] + perm_mask
    elif input_mask is not None and perm_mask is None:
      data_mask = input_mask[None]
    elif input_mask is None and perm_mask is not None:
      data_mask = perm_mask
    else:
      data_mask = None

    if data_mask is not None:
      # all mems can be attended to
      mems_mask = tf.zeros([tf.shape(data_mask)[0], mlen, bsz],
                           dtype=tf_float)
      data_mask = tf.concat([mems_mask, data_mask], 1)
      if attn_mask is None:
        attn_mask = data_mask[:, :, :, None]
      else:
        attn_mask += data_mask[:, :, :, None]

    if attn_mask is not None:
      attn_mask = tf.cast(attn_mask > 0, dtype=tf_float)

    if attn_mask is not None:
      non_tgt_mask = -tf.eye(qlen, dtype=tf_float)
      non_tgt_mask = tf.concat([tf.zeros([qlen, mlen], dtype=tf_float),
                                non_tgt_mask], axis=-1)
      non_tgt_mask = tf.cast((attn_mask + non_tgt_mask[:, :, None, None]) > 0,
                             dtype=tf_float)
    else:
      non_tgt_mask = None

    ##### Word embedding
    word_emb_k, lookup_table = embedding_lookup(
        x=inp_k,
        n_token=n_token,
        d_embed=d_model,
        initializer=initializer,
        use_tpu=use_tpu,
        dtype=tf_float,
        scope='word_embedding')

    if inp_q is not None:
      with tf.variable_scope('mask_emb'):
        mask_emb = tf.get_variable('mask_emb', [1, 1, d_model], dtype=tf_float)
        if target_mapping is not None:
          word_emb_q = tf.tile(mask_emb, [tf.shape(target_mapping)[0], bsz, 1])
        else:
          inp_q_ext = inp_q[:, :, None]
          word_emb_q = inp_q_ext * mask_emb + (1 - inp_q_ext) * word_emb_k
    output_h = tf.layers.dropout(word_emb_k, dropout, training=is_training)
    if inp_q is not None:
      output_g = tf.layers.dropout(word_emb_q, dropout, training=is_training)

    ##### Segment embedding
    if seg_id is not None:
      if untie_r:
        r_s_bias = tf.get_variable('r_s_bias', [n_layer, n_head, d_head],
                                   dtype=tf_float, initializer=initializer)
      else:
        # default case (tie)
        r_s_bias = tf.get_variable('r_s_bias', [n_head, d_head],
                                   dtype=tf_float, initializer=initializer)

      seg_embed = tf.get_variable('seg_embed', [n_layer, 2, n_head, d_head],
                                  dtype=tf_float, initializer=initializer)

      # Convert `seg_id` to one-hot `seg_mat`
      mem_pad = tf.zeros([mlen, bsz], dtype=tf.int32)
      cat_ids = tf.concat([mem_pad, seg_id], 0)

      # `1` indicates not in the same segment [qlen x klen x bsz]
      seg_mat = tf.cast(
          tf.logical_not(tf.equal(seg_id[:, None], cat_ids[None, :])),
          tf.int32)
      seg_mat = tf.one_hot(seg_mat, 2, dtype=tf_float)
    else:
      seg_mat = None

    ##### Positional encoding
    pos_emb = relative_positional_encoding(
        qlen, klen, d_model, clamp_len, attn_type, bi_data,
        bsz=bsz, dtype=tf_float)
    pos_emb = tf.layers.dropout(pos_emb, dropout, training=is_training)

    ##### Attention layers
    if mems is None:
      mems = [None] * n_layer

    for i in range(n_layer):
      # cache new mems
      new_mems.append(_cache_mem(output_h, mems[i], mem_len, reuse_len))

      # segment bias
      if seg_id is None:
        r_s_bias_i = None
        seg_embed_i = None
      else:
        r_s_bias_i = r_s_bias if not untie_r else r_s_bias[i]
        seg_embed_i = seg_embed[i]

      with tf.variable_scope('layer_{}'.format(i)):
        if inp_q is not None:
          output_h, output_g = two_stream_rel_attn(
              h=output_h,
              g=output_g,
              r=pos_emb,
              r_w_bias=r_w_bias if not untie_r else r_w_bias[i],
              r_r_bias=r_r_bias if not untie_r else r_r_bias[i],
              seg_mat=seg_mat,
              r_s_bias=r_s_bias_i,
              seg_embed=seg_embed_i,
              attn_mask_h=non_tgt_mask,
              attn_mask_g=attn_mask,
              mems=mems[i],
              target_mapping=target_mapping,
              d_model=d_model,
              n_head=n_head,
              d_head=d_head,
              dropout=dropout,
              dropatt=dropatt,
              is_training=is_training,
              kernel_initializer=initializer)
          reuse = True
        else:
          reuse = False

          output_h = rel_multihead_attn(
              h=output_h,
              r=pos_emb,
              r_w_bias=r_w_bias if not untie_r else r_w_bias[i],
              r_r_bias=r_r_bias if not untie_r else r_r_bias[i],
              seg_mat=seg_mat,
              r_s_bias=r_s_bias_i,
              seg_embed=seg_embed_i,
              attn_mask=non_tgt_mask,
              mems=mems[i],
              d_model=d_model,
              n_head=n_head,
              d_head=d_head,
              dropout=dropout,
              dropatt=dropatt,
              is_training=is_training,
              kernel_initializer=initializer,
              reuse=reuse)

        if inp_q is not None:
          output_g = positionwise_ffn(
              inp=output_g,
              d_model=d_model,
              d_inner=d_inner,
              dropout=dropout,
              kernel_initializer=initializer,
              activation_type=ff_activation,
              is_training=is_training)

        output_h = positionwise_ffn(
            inp=output_h,
            d_model=d_model,
            d_inner=d_inner,
            dropout=dropout,
            kernel_initializer=initializer,
            activation_type=ff_activation,
            is_training=is_training,
            reuse=reuse)

    if inp_q is not None:
      output = tf.layers.dropout(output_g, dropout, training=is_training)
    else:
      output = tf.layers.dropout(output_h, dropout, training=is_training)

    return output, new_mems, lookup_table


def lm_loss(hidden, target, n_token, d_model, initializer, lookup_table=None,
            tie_weight=False, bi_data=True, use_tpu=False):
  """doc."""

  with tf.variable_scope('lm_loss'):
    if tie_weight:
      assert lookup_table is not None, \
          'lookup_table cannot be None for tie_weight'
      softmax_w = lookup_table
    else:
      softmax_w = tf.get_variable('weight', [n_token, d_model],
                                  dtype=hidden.dtype, initializer=initializer)

    softmax_b = tf.get_variable('bias', [n_token], dtype=hidden.dtype,
                                initializer=tf.zeros_initializer())

    logits = tf.einsum('ibd,nd->ibn', hidden, softmax_w) + softmax_b

    if use_tpu:
      one_hot_target = tf.one_hot(target, n_token, dtype=logits.dtype)
      loss = -tf.reduce_sum(tf.nn.log_softmax(logits) * one_hot_target, -1)
    else:
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=target,
                                                            logits=logits)

    return loss


def summarize_sequence(summary_type, hidden, d_model, n_head, d_head, dropout,
                       dropatt, input_mask, is_training, initializer,
                       scope=None, reuse=None, use_proj=True):

  """
      Different classification tasks may not may not share the same parameters
      to summarize the sequence features.

      If shared, one can keep the `scope` to the default value `None`.
      Otherwise, one should specify a different `scope` for each task.
  """

  with tf.variable_scope(scope, 'sequnece_summary', reuse=reuse):
    if summary_type == 'last':
      summary = hidden[-1]
    elif summary_type == 'first':
      summary = hidden[0]
    elif summary_type == 'mean':
      summary = tf.reduce_mean(hidden, axis=0)
    elif summary_type == 'attn':
      bsz = tf.shape(hidden)[1]

      summary_bias = tf.get_variable('summary_bias', [d_model],
                                     dtype=hidden.dtype,
                                     initializer=initializer)
      summary_bias = tf.tile(summary_bias[None, None], [1, bsz, 1])

      if input_mask is not None:
        input_mask = input_mask[None, :, :, None]

      summary = multihead_attn(summary_bias, hidden, hidden, input_mask,
                               d_model, n_head, d_head, dropout, dropatt,
                               is_training, initializer, residual=False)
      summary = summary[0]
    else:
      raise ValueError('Unsupported summary type {}'.format(summary_type))

    # use another projection as in BERT
    if use_proj:
      summary = tf.layers.dense(
          summary,
          d_model,
          activation=tf.tanh,
          kernel_initializer=initializer,
          name='summary')

    # dropout
    summary = tf.layers.dropout(
        summary, dropout, training=is_training,
        name='dropout')

  return summary


def classification_loss(hidden, labels, n_class, initializer, scope, reuse=None,
                        return_logits=False):
  """
      Different classification tasks should use different scope names to ensure
      different dense layers (parameters) are used to produce the logits.

      An exception will be in transfer learning, where one hopes to transfer
      the classification weights.
  """

  with tf.variable_scope(scope, reuse=reuse):
    logits = tf.layers.dense(
        hidden,
        n_class,
        kernel_initializer=initializer,
        name='logit')

    one_hot_target = tf.one_hot(labels, n_class, dtype=hidden.dtype)
    loss = -tf.reduce_sum(tf.nn.log_softmax(logits) * one_hot_target, -1)

    if return_logits:
      return loss, logits

    return loss


def regression_loss(hidden, labels, initializer, scope, reuse=None,
                    return_logits=False):
  with tf.variable_scope(scope, reuse=reuse):
    logits = tf.layers.dense(
        hidden,
        1,
        kernel_initializer=initializer,
        name='logit')

    logits = tf.squeeze(logits, axis=-1)
    loss = tf.square(logits - labels)

    if return_logits:
      return loss, logits

    return loss