jax

db464b3f - Clarify documentation for output_offsets operand of ragged_all_to_all.