def forward(self, input_embeds: torch.Tensor, past_key_values_length: int = 0) -> torch.Tensor:
        """`input_ids_shape` is expected to be [bsz x seqlen]."""
        batch_size, seq_len = input_embeds.shape[:2]
        positions = torch.arange(
            past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=input_embeds.device
        )[:, None]
        pe = torch.zeros((seq_len, self.embed_dim), device=input_embeds.device, dtype=input_embeds.dtype)
        half_d_feature = self.embed_dim // 2
        div_term = torch.exp(
            torch.arange(half_d_feature, device=input_embeds.device, dtype=input_embeds.dtype)
            * -(np.log(float(self.max_scale)) / (half_d_feature - 1))
        )
        pe[:, :half_d_feature] = torch.sin(positions * div_term)
        pe[:, half_d_feature:] = torch.cos(positions * div_term)
        return pe[None].expand(batch_size, -1, -1)
`input_ids_shape` is expected to be [bsz x seqlen].
dpkd/transformers/src/transformers/models/pegasus_x/modeling_pegasus_x.py