Add missing fields to KJT's PyTree flatten/unflatten logic for VBE KJT#2952

Open

jd7-tr wants to merge 3 commits intopytorch:mainfrom

jd7-tr:export-D74295924

Commits on May 7, 2025

Update KJT stride calculation logic to be based off of inverse_indices for VBE KJTs. (#2949)

jd7-tr
authored and
facebook-github-bot
committed
Create a new tensor arg for stride_per_key_per_rank to facilitate torch.export (#2950)

jd7-tr
authored and
facebook-github-bot
committed
Add missing fields to KJT's PyTree flatten/unflatten logic for VBE KJT

jd7-tr
authored and
facebook-github-bot
committed