import tensorflow as tf

                )

                self_attn_input = rep_map
                context_features = tf.add(scatter_attn[:, :-1], scatter_pooling[:, :-1], 'context_features')
                output_mask = rep_mask
            else:
                self_attn_input = rep_head_tensor
                context_features = attn_result
                output_mask = rep_head_mask

            # context fusion gate
            o_bias = tf.get_variable('o_bias', [ivec], tf.float32, tf.constant_initializer(0.))
            fusion_gate = tf.nn.sigmoid(
                linear(self_attn_input, ivec, True, 0., 'linear_fusion_i', False, wd, keep_prob, is_train) +
                linear(context_features, ivec, True, 0., 'linear_fusion_a', False, wd, keep_prob, is_train) +
                o_bias)
            output = fusion_gate * self_attn_input + (1 - fusion_gate) * context_features

        return output, output_mask


def self_attention_for_selected_head(
