import tensorflow as tf

        # non-linear
        rep_map = bn_dense_layer(rep_tensor, ivec, True, 0., 'bn_dense_map', activation,
                                 False, wd, keep_prob, is_train)
        # ensure the seletion is right
        dep_selection = tf.logical_and(rep_mask, dep_selection)
        head_selection = tf.logical_and(rep_mask, head_selection)
        rep_dep_tensor, rep_dep_mask, dep_org_idx = reduce_data_rep_max_len(rep_map, dep_selection)
        rep_head_tensor,rep_head_mask, head_org_idx = reduce_data_rep_max_len(rep_map, head_selection)
        sl_dep, sl_head = tf.shape(rep_dep_tensor)[1], tf.shape(rep_head_tensor)[1]

        if keep_unselected:
            unhead_selection = tf.logical_and(rep_mask, tf.logical_not(head_selection))
            rep_unhead_tensor, rep_unhead_mask, unhead_org_idx = reduce_data_rep_max_len(rep_map, unhead_selection)
            sl_unhead = tf.shape(rep_unhead_tensor)[1]

        attn_result = tf.cond(
            tf.equal(sl_head, 0),
            lambda: tf.zeros([bs, 0, hn], tf.float32),
            lambda: self_attention_for_selected_head(
                head_selection, head_org_idx, sl_head, rep_head_mask,
                dep_selection, dep_org_idx, sl_dep, rep_dep_mask,
                rep_map, rep_dep_tensor, keep_prob, is_train, direction, ivec
            )
        )

        if keep_unselected:
