(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)
return attn_vec
Core relative positional attention operations.
Core relative positional attention operations.
Core relative positional attention operations .
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_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 ) : # 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