jax.lax.pmax ============ .. currentmodule:: jax.lax .. autofunction:: pmax