jax

5afc7758 - Merge pull request #4868 from sharadmv:variadic_reduce