import tensorflow as tf

                        tf.stack([range_head, head_org_idx], -1), attn_result, [bs, sl+1, hn])
                )

                range_unhead = tf.tile(tf.expand_dims(tf.range(bs), -1), [1, sl_unhead])
                scatter_pooling = tf.cond(
                    tf.equal(sl_unhead, 0),
                    lambda: tf.zeros([bs, sl+1, hn], tf.float32),
                    lambda: tf.scatter_nd(
                        tf.stack([range_unhead, unhead_org_idx], -1), pooling_result, [bs, sl+1, hn])
                )

                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
