[JAX] Fix get_seqlens_and_offsets() to accept vmapped seg ids and non vmapped seg offsets#2692
Open
KshitijLakhani wants to merge 12 commits intoNVIDIA:mainfrom
Open
[JAX] Fix get_seqlens_and_offsets() to accept vmapped seg ids and non vmapped seg offsets#2692KshitijLakhani wants to merge 12 commits intoNVIDIA:mainfrom
KshitijLakhani wants to merge 12 commits intoNVIDIA:mainfrom
Commits
Commits on Feb 24, 2026
![pre-commit-ci[bot]](https://avatars.githubusercontent.com/in/68672?v=4&size=32){kind=avatar}
Commits on Feb 25, 2026
Commits on Feb 27, 2026
- committed
- committed
Undo batcher check logic for seg pos and seg ids as it is already moved to get_seqlens_and_offsets()
committed- committed
- committed
- committed
- authored