jax

141fabbb - Reimplement argmin/argmax using a single pass variadic reduction. (#3611)