jax

126c5b31 - Merge pull request #4736 from bchetioui:reduce_scatter_tests