jax.lax.custom\_root ==================== .. currentmodule:: jax.lax .. autofunction:: custom_root