jax.lax.shift_right_logical