def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False):
    batch_size = shape_list(hidden_states)[0]
    mixed_query_layer = self.query(hidden_states)
    mixed_key_layer = self.key(hidden_states)
    mixed_value_layer = self.value(hidden_states)
    mixed_key_conv_attn_layer = self.key_conv_attn_layer(hidden_states)
    query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
    key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
    conv_attn_layer = tf.multiply(mixed_key_conv_attn_layer, mixed_query_layer)
    conv_kernel_layer = self.conv_kernel_layer(conv_attn_layer)
    conv_kernel_layer = tf.reshape(conv_kernel_layer, [(- 1), self.conv_kernel_size, 1])
    conv_kernel_layer = tf.nn.softmax(conv_kernel_layer, axis=1)
    paddings = tf.constant([[0, 0], [int(((self.conv_kernel_size - 1) / 2)), int(((self.conv_kernel_size - 1) / 2))], [0, 0]])
    conv_out_layer = self.conv_out_layer(hidden_states)
    conv_out_layer = tf.reshape(conv_out_layer, [batch_size, (- 1), self.all_head_size])
    conv_out_layer = tf.pad(conv_out_layer, paddings, 'CONSTANT')
    unfold_conv_out_layer = tf.stack([tf.slice(conv_out_layer, [0, i, 0], [batch_size, shape_list(mixed_query_layer)[1], self.all_head_size]) for i in range(self.conv_kernel_size)], axis=(- 1))
    conv_out_layer = tf.reshape(unfold_conv_out_layer, [(- 1), self.attention_head_size, self.conv_kernel_size])
    conv_out_layer = tf.matmul(conv_out_layer, conv_kernel_layer)
    conv_out_layer = tf.reshape(conv_out_layer, [(- 1), self.all_head_size])
    attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
    dk = tf.cast(shape_list(key_layer)[(- 1)], attention_scores.dtype)
    attention_scores = (attention_scores / tf.math.sqrt(dk))
    if (attention_mask is not None):
        attention_scores = (attention_scores + attention_mask)
    attention_probs = tf.nn.softmax(attention_scores, axis=(- 1))
    attention_probs = self.dropout(attention_probs, training=training)
    if (head_mask is not None):
        attention_probs = (attention_probs * head_mask)
    value_layer = tf.reshape(mixed_value_layer, [batch_size, (- 1), self.num_attention_heads, self.attention_head_size])
    value_layer = tf.transpose(value_layer, [0, 2, 1, 3])
    context_layer = tf.matmul(attention_probs, value_layer)
    context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
    conv_out = tf.reshape(conv_out_layer, [batch_size, (- 1), self.num_attention_heads, self.attention_head_size])
    context_layer = tf.concat([context_layer, conv_out], 2)
    context_layer = tf.reshape(context_layer, (batch_size, (- 1), (self.head_ratio * self.all_head_size)))
    outputs = ((context_layer, attention_probs) if output_attentions else (context_layer,))
    return outputs
def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False):
    batch_size = shape_list(hidden_states)[0]
    mixed_query_layer = self.query(hidden_states)
    mixed_key_layer = self.key(hidden_states)
    mixed_value_layer = self.value(hidden_states)
    mixed_key_conv_attn_layer = self.key_conv_attn_layer(hidden_states)
    query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
    key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
    conv_attn_layer = tf.multiply(mixed_key_conv_attn_layer, mixed_query_layer)
    conv_kernel_layer = self.conv_kernel_layer(conv_attn_layer)
    conv_kernel_layer = tf.reshape(conv_kernel_layer, [(- 1), self.conv_kernel_size, 1])
    conv_kernel_layer = tf.nn.softmax(conv_kernel_layer, axis=1)
    paddings = tf.constant([[0, 0], [int(((self.conv_kernel_size - 1) / 2)), int(((self.conv_kernel_size - 1) / 2))], [0, 0]])
    conv_out_layer = self.conv_out_layer(hidden_states)
    conv_out_layer = tf.reshape(conv_out_layer, [batch_size, (- 1), self.all_head_size])
    conv_out_layer = tf.pad(conv_out_layer, paddings, 'CONSTANT')
    unfold_conv_out_layer = tf.stack([tf.slice(conv_out_layer, [0, i, 0], [batch_size, shape_list(mixed_query_layer)[1], self.all_head_size]) for i in range(self.conv_kernel_size)], axis=(- 1))
    conv_out_layer = tf.reshape(unfold_conv_out_layer, [(- 1), self.attention_head_size, self.conv_kernel_size])
    conv_out_layer = tf.matmul(conv_out_layer, conv_kernel_layer)
    conv_out_layer = tf.reshape(conv_out_layer, [(- 1), self.all_head_size])
    attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
    dk = tf.cast(shape_list(key_layer)[(- 1)], attention_scores.dtype)
    attention_scores = (attention_scores / tf.math.sqrt(dk))
    if (attention_mask is not None):
        attention_scores = (attention_scores + attention_mask)
    attention_probs = tf.nn.softmax(attention_scores, axis=(- 1))
    attention_probs = self.dropout(attention_probs, training=training)
    if (head_mask is not None):
        attention_probs = (attention_probs * head_mask)
    value_layer = tf.reshape(mixed_value_layer, [batch_size, (- 1), self.num_attention_heads, self.attention_head_size])
    value_layer = tf.transpose(value_layer, [0, 2, 1, 3])
    context_layer = tf.matmul(attention_probs, value_layer)
    context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
    conv_out = tf.reshape(conv_out_layer, [batch_size, (- 1), self.num_attention_heads, self.attention_head_size])
    context_layer = tf.concat([context_layer, conv_out], 2)
    context_layer = tf.reshape(context_layer, (batch_size, (- 1), (self.head_ratio * self.all_head_size)))
    outputs = ((context_layer, attention_probs) if output_attentions else (context_layer,))
    return outputs