jax.lax.switch ============== .. currentmodule:: jax.lax .. autofunction:: switch