def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0):
        """Input is expected to be of size [bsz x seqlen]."""
        bsz, seq_len = input_shape[:2]

        positions = tf.range(past_key_values_length, seq_len + past_key_values_length, delta=1, name="range")
        return tf.gather(self.weight, positions)
Input is expected to be of size [bsz x seqlen].
dpkd/transformers/src/transformers/models/roformer/modeling_tf_roformer.py