[JAX] Fix incorrect calculation of segment pos from segment ids in user-facing API#2523

Merged

ksivaman merged 19 commits intoNVIDIA:mainfrom

KshitijLakhani:klakhani/fix/incorrect-sequence-descr-from-seg-ids-only

Dec 31, 2025

Commits on Dec 22, 2025

Fix incorrect calculation of segment pos from segment ids for thd cases and load balanced cases in from_segment_ids_and_pos. Enforce passing of segment_pos for THD cases and lod balanced cases

KshitijLakhani
authored and
Kshitij Janardan Lakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
Kshitij Janardan Lakhani
committed
Correct the assert condition

KshitijLakhani
authored and
Kshitij Janardan Lakhani
committed
Modify fused attn tests to pass new args to from_segment_ids_and_pos()

KshitijLakhani
authored and
Kshitij Janardan Lakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
Kshitij Janardan Lakhani
committed
Calculate seg ids before pos

KshitijLakhani
authored and
Kshitij Janardan Lakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
Kshitij Janardan Lakhani
committed

Commits on Dec 23, 2025

1. Change the signature for from_segment_ids_and_pos()
Kshitij Janardan Lakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed
Pass keyword-only args by name
Kshitij Janardan Lakhani
committed
nit: Fix typo to use seg_ids instead of segment_ids
Kshitij Janardan Lakhani
committed
nit: Fix comments
Kshitij Janardan Lakhani
committed
Modify the function call to differentiate between load balancing and actually reordered segment_ids and segment_pos
Kshitij Janardan Lakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed

Commits on Dec 31, 2025

Merge branch 'main' into klakhani/fix/incorrect-sequence-descr-from-seg-ids-only
ksivaman
authored