import tensorflow as tf


        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:
            input_idx = tf.tile(tf.expand_dims(tf.range(sl), 0), [bs, 1])
            pooling_result = tf.cond(
                tf.equal(sl_unhead, 0),
                lambda: tf.zeros([bs, 0, hn], tf.float32),
                lambda: mean_pooling_for_unselected_head(
                    unhead_org_idx, sl_unhead, rep_unhead_mask,
                    input_idx, sl, rep_mask, rep_map, None)  # todo: point !
            )

        with tf.variable_scope('output'):
            if keep_unselected:
                range_head = tf.tile(tf.expand_dims(tf.range(bs), -1), [1, sl_head])
                scatter_attn = tf.cond(
                    tf.equal(sl_head, 0),
                    lambda: tf.zeros([bs, sl+1, hn], tf.float32),
                    lambda: tf.scatter_nd(
