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

        return super().call(tf.cast(position_ids, dtype=tf.int32))
Input is expected to be of size [bsz x seqlen].
dpkd/transformers/src/transformers/models/led/modeling_tf_led.py