src/transformers/models/transfo_xl/modeling_tf_transfo_xl.py/TFRelPartialLearnableMultiHeadAttn/_rel_shift
class TFRelPartialLearnableMultiHeadAttn:    def _rel_shift(self, x):
        x_size = shape_list(x)

        x = tf.pad(x, [[0, 0], [1, 0], [0, 0], [0, 0]])
        x = tf.reshape(x, [x_size[1] + 1, x_size[0], x_size[2], x_size[3]])
        x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
        x = tf.reshape(x, x_size)

        return x