(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.
Two-stream attention with relative positional encoding.
Two - stream attention with relative positional encoding .
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 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' ) : 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