jax

80973185 - cleanup: unify pmin/pmax implementations with psum