code: def rel_multihead_attn(
      w,
      r,
      r_w_bias,
      r_r_bias,
      attn_mask,
      mems,
      d_model,
      n_head,
      d_head,
      dropout,
      dropatt,
      is_training,
      kernel_initializer,
      index,
      scope="rel_attn",
  ):
      scale = 1 / (d_head**0.5)
      with tf.variable_scope(scope):
          qlen = tf.shape(w)[0]
          rlen = tf.shape(r)[0]
          bsz = tf.shape(w)[1]
  
          cat = (
              tf.concat([mems, w], 0)
              if mems is not None and mems.shape.ndims > 1
              else w
          )
          # w_heads = tf.layers.dense(cat, 3 * n_head * d_head, use_bias=False,
          # kernel_initializer=kernel_initializer, name='qkv')
          if index <= 3:
              rank = 75
          else:
              rank = 75
          w_heads = tf.layers.dense(
              cat,
              3 * rank,
              use_bias=False,
              kernel_initializer=kernel_initializer,
              name="qkv_",
          )
          tmp = w_heads
          w_heads = tf.layers.dense(
              w_heads,
              3 * n_head * d_head,
              use_bias=False,
              kernel_initializer=kernel_initializer,
              name="qkv",
          )
          r_head_k = tf.layers.dense(
              r,
              n_head * d_head,
              use_bias=False,
              kernel_initializer=kernel_initializer,
              name="r",
          )
  
          w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, -1)
          w_head_q = w_head_q[-qlen:]
  
          klen = tf.shape(w_head_k)[0]
  
          w_head_q = tf.reshape(w_head_q, [qlen, bsz, n_head, d_head])
          w_head_k = tf.reshape(w_head_k, [klen, bsz, n_head, d_head])
          w_head_v = tf.reshape(w_head_v, [klen, bsz, n_head, d_head])
  
          r_head_k = tf.reshape(r_head_k, [rlen, n_head, d_head])
  
          rw_head_q = w_head_q + r_w_bias
          rr_head_q = w_head_q + r_r_bias
  
          AC = tf.einsum("ibnd,jbnd->ijbn", rw_head_q, w_head_k)
          BD = tf.einsum("ibnd,jnd->ijbn", rr_head_q, r_head_k)
          bd_before = BD
          BD = rel_shift(BD)
  
          attn_score = (AC + BD) * scale
          attn_mask_t = attn_mask[:, :, None, None]
          attn_score = (
              attn_score * (1 - attn_mask_t) - 1e30 * attn_mask_t
          )
  
          attn_prob = tf.nn.softmax(attn_score, 1)
          attn_prob = tf.layers.dropout(
              attn_prob, dropatt, training=is_training
          )
  
          attn_vec = tf.einsum("ijbn,jbnd->ibnd", attn_prob, w_head_v)
          size_t = tf.shape(attn_vec)
          attn_vec = tf.reshape(
              attn_vec, [size_t[0], size_t[1], n_head * d_head]
          )
  
          attn_out = tf.layers.dense(
              attn_vec,
              d_model,
              use_bias=False,
              kernel_initializer=kernel_initializer,
              name="o",
          )
          attn_out = tf.layers.dropout(
              attn_out, dropout, training=is_training
          )
          befornorm = attn_out
  
          output = tf.contrib.layers.layer_norm(
              attn_out + w, begin_norm_axis=-1
          )
      return output
  
file_path: delta/gcompiler/example/python/nlp_transformer/model.py
imports: import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: rel_multihead_attn(w, r, r_w_bias, r_r_bias, attn_mask, mems, d_model, n_head, d_head, dropout, dropatt, is_training, kernel_initializer, index, scope) and returns: output