Expose jax.lax.all_gather (#2449)

Commit

5 years ago

Expose jax.lax.all_gather (#2449) * Expose jax.lax.all_gather * add all_gather to RTD