import tensorflow as tf

            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(
