How does jax implement jax.lax.psum and jax.lax.all_to_all?

[anirudhitagi]

Hello all,

I am new to Jax and I am exploring collective communications and so I wish to understand how does Jax implement jax.lax.psum, jax.lax.all_gather and jax.lax.all_to_all ?

Is it implemented as Ring, Binary tree, halving-doubling or something else?
And does this implementation depend on the hardware it is being deployed on? For context, I am experimenting with TPU v-3s.

Any documentation or explanation would be most appreciated !

Thanks and Regards,
AI