(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.
Defines a Transformer-XL computation graph with additional
    support for XLNet.
Defines a Transformer - XL computation graph with additional support for XLNet .
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 = inp_k.get_shape()[1] #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 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 ) : 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 = inp_k . get_shape ( ) [ 1 ] #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