Update the _maybe_compute_stride_kjt logic to calculate stride based off of inverse_indices for VBE KJTs.#2925

Closed

jd7-tr wants to merge 2 commits intopytorch:mainfrom

jd7-tr:export-D73824764

Commits on May 1, 2025

Include stride_per_key_per_rank in KJT's PyTree flatten/unflatten logic (#2903)

James Dong
authored and
facebook-github-bot
committed
Update the _maybe_compute_stride_kjt logic to calculate stride based off of inverse_indices for VBE KJTs. (#2925)

jd7-tr
authored and
facebook-github-bot
committed