jax

8f29cb71 - Merge pull request #40572 from aybchan:hlo-collective-compat