jax.lax.psum ============ .. currentmodule:: jax.lax .. autofunction:: psum