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

ksivaman merged 19 commits into

NVIDIA:mainNVIDIA/TransformerEngine:mainfrom

KshitijLakhani:klakhani/fix/incorrect-sequence-descr-from-seg-ids-onlyKshitijLakhani/TransformerEngine:klakhani/fix/incorrect-sequence-descr-from-seg-ids-onlyCopy head branch name to clipboard