Skip to content

[tx] replace calls to jax.process_index() to resolve rank ordering issue with multi-host TPUs#1252

Merged
pcmoritz merged 11 commits intoNovaSky-AI:mainfrom
andrewsykim:remove-jax-process-id
Mar 4, 2026
Merged

[tx] replace calls to jax.process_index() to resolve rank ordering issue with multi-host TPUs#1252
pcmoritz merged 11 commits intoNovaSky-AI:mainfrom
andrewsykim:remove-jax-process-id

Commits

Commits on Mar 2, 2026

Commits on Mar 3, 2026