import tensorflow as tf

            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(
