The PyTorch API reference for flashinfer.rope.apply_rope_pos_ids appears to contain inaccuracies.

[ovowei]

In the code at this line, the pos_ids used have a shape of (nnz). However, the documentation at flashinfer.rope.apply_rope_pos_ids describes pos_ids as having a shape of (batch_size + 1).

[yzh119]

Thank you for reporting this issue, pos_ids should have shape (nnz,), and the documentation was not updated.