Comprehensive User Manual and Reference

Guide
JAX Nonlinear System Solvers
High-Performance Newton-Raphson, Trust-Region, and Broyden Methods

July 25, 2025

Abstract
This comprehensive manual provides detailed documentation for high-performance, JAX-compatible solvers
for nonlinear systems of equations F (x) = 0 in Rn . The implementation leverages JAX’s automatic
differentiation and just-in-time compilation capabilities to achieve both numerical robustness and com-
putational efficiency. This document serves as both a user guide and a theoretical reference, covering
Newton-Raphson methods with Strong Wolfe line search, trust-region methods with dog-leg steps, and
hybrid Broyden methods with automatic fallback strategies.
Contents
1 Introduction 4
1.1 Key Features . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4
1.2 When to Use This Library . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4

2 Mathematical Background 4
2.1 The Nonlinear System Problem . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4
2.2 Newton-Raphson Method . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4
2.3 Globalization Strategies . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5
2.3.1 Line Search Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5
2.3.2 Trust-Region Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5
2.4 Broyden’s Method . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5

3 Algorithms Overview 5
3.1 Newton with Line Search . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5
3.2 Trust-Region with Dog-leg . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5
3.3 Hybrid Broyden . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5

4 Installation and Setup 7

4.1 Requirements . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7
4.2 Installation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7
4.3 Configuration . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7

5 API Reference 7
5.1 newton() . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7
5.2 broyden_hybrid() . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8

6 Usage Examples 8
6.1 Basic Example: Finding Square Root . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8
6.2 Multidimensional System . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8
6.3 Stiff System Example . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9
6.4 Large-Scale Problem . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9

7 Performance Tuning Guide 9

7.1 Parameter Selection Guidelines . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9
7.2 Method Selection . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9
7.2.1 When to Use Newton with Line Search . . . . . . . . . . . . . . . . . . . . . . . . . . . 9
7.2.2 When to Use Trust Region . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10
7.2.3 When to Use Broyden . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10
7.3 Performance Optimization Tips . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10

8 Troubleshooting Guide 10
8.1 Common Issues and Solutions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10
8.1.1 Issue: Solver reaches maximum iterations . . . . . . . . . . . . . . . . . . . . . . . . . 10
8.1.2 Issue: NaN or Inf in results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11
8.1.3 Issue: Slow convergence . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11
8.2 Debugging Strategies . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11

9 Advanced Topics 11
9.1 Handling Singular Jacobians . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11
9.2 Continuation Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12
9.3 Parallel Solution of Multiple Systems . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12

10 Theoretical Foundations 12
10.1 Convergence Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12
10.2 Trust Region Theory . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13
10.3 Line Search Theory . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13

2
JAX Nonlinear System Solvers - Reference Manual 3

10.3.1 Armijo Condition (Sufficient Decrease) . . . . . . . . . . . . . . . . . . . . . . . . . . . 13

10.3.2 Wolfe Conditions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13
10.3.3 Strong Wolfe Conditions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13
10.4 Complexity Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13

11 References and Further Reading 13

11.1 Core References . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13
11.2 Newton-Raphson Method . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14
11.3 Trust Region Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14
11.4 Broyden’s Method and Quasi-Newton Methods . . . . . . . . . . . . . . . . . . . . . . . . . . 14
11.5 Strong Wolfe Line Search . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14
11.6 JAX and Automatic Differentiation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15
11.7 Numerical Analysis and Implementation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15
11.8 Specialized Topics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15
11.9 Recent Advances . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15
11.10Online Resources . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16
11.11Software Libraries . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16

A Complete Code Example 16

B Common Patterns and Best Practices 17

B.1 Problem Scaling . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 17
B.2 Robustness Wrapper . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18
B.3 Convergence Monitoring . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18

C Troubleshooting Checklist 19
JAX Nonlinear System Solvers - Reference Manual 4

1 Introduction
This module provides high-performance, JAX-compatible solvers for nonlinear systems of equations F (x) = 0
in Rn . The implementation leverages JAX’s automatic differentiation and just-in-time compilation capabil-
ities to achieve both numerical robustness and computational efficiency.

1.1 Key Features

• Newton-Raphson Method with Strong Wolfe line search

• Trust-Region Method with dog-leg step computation

• Hybrid Broyden Method with automatic fallback to Newton steps

• Full JAX compatibility for GPU acceleration and automatic differentiation

• Robust globalization strategies for improved convergence

1.2 When to Use This Library

• Solving systems of nonlinear equations F (x) = 0

• Problems requiring high numerical precision (64-bit arithmetic)

• Large-scale systems where Jacobian computation is expensive

• Applications requiring GPU acceleration or automatic differentiation

2 Mathematical Background
2.1 The Nonlinear System Problem
Given a function F : D ⊂ Rn → Rn , we seek x∗ such that F (x∗ ) = 0. This arises in:
• Steady-state analysis of dynamical systems

• Optimization (finding critical points where ∇f (x) = 0)

• Implicit function problems

• Numerical solution of PDEs

2.2 Newton-Raphson Method

The Newton-Raphson method iteratively solves:

J(xk )∆x = −F (xk ) (1)

xk+1 = xk + ∆x (2)

where J(x) = ∂F/∂x is the Jacobian matrix.

Theorem 2.1 (Newton’s Method Convergence). Let F : D ⊂ Rn → Rn be continuously differentiable in an
open convex set D. Assume:
1. ∃x∗ ∈ D such that F (x∗ ) = 0

2. J(x∗ ) is nonsingular

3. J is Lipschitz continuous: ∥J(x) − J(y)∥ ≤ L ∥x − y∥

Then ∃δ > 0 such that for x0 with ∥x0 − x∗ ∥ < δ, Newton’s method converges quadratically:
L
∥xk+1 − x∗ ∥ ≤ ∥xk − x∗ ∥2 (3)
2 ∥J(x∗ )−1 ∥
JAX Nonlinear System Solvers - Reference Manual 5

2.3 Globalization Strategies

2.3.1 Line Search Methods
Find α ∈ (0, 1] such that:
∥F (xk + α∆x)∥ < ∥F (xk )∥ (4)

Strong Wolfe Conditions:

1. Sufficient decrease: ϕ(α) ≤ ϕ(0) + c1 αϕ′ (0)

2. Curvature condition: |ϕ′ (α)| ≤ c2 |ϕ′ (0)|

where ϕ(α) = 1
2 ∥F (xk + α∆x)∥2 .

2.3.2 Trust-Region Methods

Solve the constrained minimization:
1 1
minimizep mk (p) = ∥F (xk )∥2 + pT JkT Fk + pT JkT Jk p subject to ∥p∥ ≤ ∆ (5)
2 2

2.4 Broyden’s Method

Quasi-Newton method that approximates the Jacobian:

(yk − Bk sk )sTk
Bk+1 = Bk + (6)
∥sk ∥2

where yk = F (xk+1 ) − F (xk ) and sk = xk+1 − xk .

Theorem 2.2 (Broyden’s Method Convergence). Under similar conditions to Newton’s method, Broyden’s
method achieves superlinear convergence:
∥xk+1 − x∗ ∥
lim =0 (7)
k→∞ ∥xk − x∗ ∥

3 Algorithms Overview
3.1 Newton with Line Search

Algorithm 1 Newton with Strong Wolfe Line Search

Require: Initial guess x0 , tolerance ε, max iterations N
Ensure: Approximate solution x∗ such that ∥F (x∗ )∥ < ε
1: for k = 0, 1, . . . , N − 1 do
2: Compute F (xk ) and J(xk )
3: Solve J(xk )p = −F (xk ) for Newton direction p
4: Find αk satisfying Strong Wolfe conditions
5: Update xk+1 = xk + αk p
6: if ∥F (xk+1 )∥ < ε then
7: return xk+1
8: end if
9: end for

3.2 Trust-Region with Dog-leg

3.3 Hybrid Broyden

JAX Nonlinear System Solvers - Reference Manual 6

Algorithm 2 Dog-leg Trust Region Step

Require: Current point x, trust radius ∆
Ensure: Dog-leg step p
1: Compute Cauchy point: pU = −τ U g where τ U minimizes m along −g
2: Compute Newton point: pB = −J −1 F
3: if pB ≤ ∆ then
4: return pB
5: else if pU ≥ ∆ then
6: return (∆/ pU )pU
7: else
8: Find intersection of dog-leg path with trust region boundary
9: return intersection point
10: end if

Algorithm 3 Hybrid Broyden with Newton Fallback

Require: Initial guess x0 , initial Jacobian B0 = J(x0 )
Ensure: Approximate solution x∗
1: for k = 0, 1, . . . , N − 1 do
2: Solve Bk p = −F (xk ) for Broyden direction
3: Perform line search with step p
4: if step fails to reduce ∥F ∥ then
5: Compute Newton step as fallback
6: end if
7: Update Bk+1 using Broyden formula
8: if k mod refresh_freq = 0 then
9: Refresh Bk+1 = J(xk+1 )
10: end if
11: end for
 JAX Nonlinear System Solvers - Reference Manual 7

4 Installation and Setup

4.1 Requirements

Python >= 3.8

JAX >= 0.4.0
jaxlib ( matching JAX version )
numpy >= 1.20.0

4.2 Installation

# Install JAX with CPU support

pip install -- upgrade jax jaxlib

# For GPU support ( CUDA 11. x )

pip install -- upgrade " jax [ cuda11_pip ] " -f \
https :// storage . googleapis . com / jax - releases / jax_cuda_releases . html

# For GPU support ( CUDA 12. x )

pip install -- upgrade " jax [ cuda12_pip ] " -f \
https :// storage . googleapis . com / jax - releases / jax_cuda_releases . html

4.3 Configuration

1 import jax
2 # Enable 64 - bit precision ( recommended for numerical stability )
3 jax . config . update ( " jax_enable_x64 " , True )
4
5 # Optional : Set default device
6 jax . config . update ( " jax_default_device " , " gpu " ) # or " cpu "

5 API Reference
5.1 newton()

1 newton ( fun , x0 , maxiter =100 , tol =1 e -8 , method = ’ line_search ’ ,

2 delta0 =1.0 , alpha0 =1.0 , c1 =1 e -4 , c2 =0.9 , rho =0.5)

Parameters:
• fun: Callable Rn → Rn - The nonlinear function F (x)
• x0: jnp.ndarray - Initial guess
• maxiter: int - Maximum iterations (default: 100)
• tol: float - Convergence tolerance on ∥F (x)∥ (default: 1e-8)
• method: str - ’line_search’ or ’trust_region’ (default: ’line_search’)
• delta0: float - Initial trust region radius (default: 1.0)
• alpha0: float - Initial line search step (default: 1.0)
• c1: float - Wolfe sufficient decrease parameter (default: 1e-4)
• c2: float - Wolfe curvature parameter (default: 0.9)
• rho: float - Line search backtracking factor (default: 0.5)
Returns:
• x_root: jnp.ndarray - Approximate solution
• n_iters: int - Number of iterations performed
 JAX Nonlinear System Solvers - Reference Manual 8

5.2 broyden_hybrid()

1 broyden_hybrid ( fun , x0 , B0 = None , maxiter =100 , tol =1 e -8 ,

2 refresh_freq =15 , alpha0 =1.0 , c1 =1 e -4 , c2 =0.9 , rho =0.5)

Parameters:

• fun: Callable Rn → Rn - The nonlinear function F (x)

• x0: jnp.ndarray - Initial guess

• B0: jnp.ndarray or None - Initial Jacobian approximation (default: J(x0 ))

• maxiter: int - Maximum iterations (default: 100)

• tol: float - Convergence tolerance on ∥F (x)∥ (default: 1e-8)

• refresh_freq: int - Jacobian refresh frequency (default: 15)

• Other parameters same as newton()

6 Usage Examples
6.1 Basic Example: Finding Square Root

1 import jax
2 import jax . numpy as jnp
3 from nonlin ear_solvers import newton , broyden_hybrid
4
5 # Enable 64 - bit precision
6 jax . config . update ( " jax_enable_x64 " , True )
7
8 # Define system : x ^2 - 2 = 0
9 def f ( x ) :
10 return jnp . array ([ x [0]**2 - 2.0])
11
12 # Initial guess
13 x0 = jnp . array ([1.0])
14
15 # Solve with Newton
16 x_newton , iters_newton = newton (f , x0 )
17 print ( f " Newton : x = { x_newton [0]:.12 f } , iterations = { iters_newton } " )
18
19 # Solve with Broyden
20 x_broyden , iters_broyden = broyden_hybrid (f , x0 )
21 print ( f " Broyden : x = { x_broyden [0]:.12 f } , iterations = { iters_broyden } " )

6.2 Multidimensional System

1 # System of equations :
2 # x ^2 + y ^2 = 4
3 # x - y = 0
4 def system_2d ( x ) :
5 return jnp . array ([
6 x [0]**2 + x [1]**2 - 4.0 ,
7 x [0] - x [1]
8 ])
9
10 x0 = jnp . array ([2.0 , 1.0])
11 x_sol , iters = newton ( system_2d , x0 )
12 print ( f " Solution : x = { x_sol [0]:.6 f } , y = { x_sol [1]:.6 f } " )
13 print ( f " Verification : F ( x ) = { system_2d ( x_sol ) } " )
 JAX Nonlinear System Solvers - Reference Manual 9

6.3 Stiff System Example

1 # Robertson chemical kinetics ( steady state )

2 def robertson ( x ) :
3 return jnp . array ([
4 -0.04* x [0] + 1 e4 * x [1]* x [2] ,
5 0.04* x [0] - 1 e4 * x [1]* x [2] - 3 e7 * x [1]**2 ,
6 x [0] + x [1] + x [2] - 1.0
7 ])
8
9 # This system is stiff - use trust region method
10 x0 = jnp . array ([1.0 , 0.0 , 0.0])
11 x_sol , iters = newton ( robertson , x0 , method = ’ trust_region ’)
12 print ( f " Solution : { x_sol } " )
13 print ( f " Residual norm : { jnp . linalg . norm ( robertson ( x_sol ) ) } " )

6.4 Large-Scale Problem

1 # Nonlinear system : x_i ^2 + sum ( x ) - 2 N = 0

2 def large_system (x , N =100) :
3 return x **2 + jnp . sum ( x ) - 2* N
4
5 # Use Broyden for efficiency with large systems
6 N = 100
7 x0 = jnp . ones ( N ) * 0.5
8 x_sol , iters = broyden_hybrid ( large_system , x0 , refresh_freq =10)
9 print ( f " Solution norm : { jnp . linalg . norm ( x_sol ) } " )
10 print ( f " Iterations : { iters } " )

7 Performance Tuning Guide

7.1 Parameter Selection Guidelines

Table 1: Recommended Parameter Values

Parameter Default Conservative Aggressive

tol 10−8 10−12 to 10−10 10−6 to 10−4
maxiter 100 200-500 50-100
alpha0 1.0 0.1-0.5 1.0-2.0
c1 10−4 10−5 10−3
c2 0.9 0.1-0.5 0.9-0.99
rho 0.5 0.8-0.9 0.1-0.3
refresh_freq 15 5-10 20-50

7.2 Method Selection

7.2.1 When to Use Newton with Line Search
• Small to medium systems (n < 1000)

• High accuracy requirements

• Well-conditioned problems

• When Jacobian is cheap to compute

JAX Nonlinear System Solvers - Reference Manual 10

7.2.2 When to Use Trust Region

• Highly nonlinear problems

• Poor initial guesses

• Problems with multiple solutions

• When line search fails frequently

7.2.3 When to Use Broyden

• Large systems (n > 1000)

• Expensive Jacobian computation

• Moderate accuracy requirements

• Problems with slowly-varying Jacobians

7.3 Performance Optimization Tips

1. Use JIT Compilation
1 solver_jit = jax . jit ( newton , static_argnums =(0 ,) )

2. Batch Processing
1 # Solve multiple systems in parallel
2 vmap_newton = jax . vmap ( newton , in_axes =( None , 0) )
3 x0_batch = jnp . stack ([ x0_1 , x0_2 , ... , x0_n ])
4 solutions = vmap_newton ( fun , x0_batch )

3. GPU Acceleration
1 # Ensure operations run on GPU
2 with jax . default_device ( jax . devices ( ’ gpu ’) [0]) :
3 x_sol , iters = newton ( fun , x0 )

8 Troubleshooting Guide
8.1 Common Issues and Solutions
8.1.1 Issue: Solver reaches maximum iterations
Symptoms: n_iters == maxiter, large residual norm
Solutions:

1. Increase maxiter (try 2x or 5x current value)

2. Relax tolerance (tol = 1e-6 instead of 1e-8)

3. Improve initial guess using domain knowledge

4. Switch methods (try trust region if using line search)

5. Check problem scaling - consider normalizing variables

 JAX Nonlinear System Solvers - Reference Manual 11

8.1.2 Issue: NaN or Inf in results

Symptoms: Solution contains NaN/Inf, crashes during iteration
Solutions:
1. Decrease alpha0 (try 0.1 or 0.01)

2. Use trust region method for better stability

3. Check function implementation for division by zero

4. Enable numerical checks:

1 jax . config . update ( " jax_debug_nans " , True )

5. Add regularization to ill-conditioned Jacobians:

1 J_reg = J + 1e -8 * jnp . eye ( n )

8.1.3 Issue: Slow convergence

Symptoms: Many iterations, slow residual decrease
Solutions:
1. For Broyden: decrease refresh_freq (more Jacobian updates)

2. Adjust line search parameters:

• Increase c1 for easier step acceptance

• Decrease c2 for stricter curvature

3. Try different initial guess strategies

4. Consider problem preconditioning

8.2 Debugging Strategies

1 def newton_debug ( fun , x0 , ** kwargs ) :

2 x = x0
3 for i in range ( kwargs . get ( ’ maxiter ’ , 100) ) :
4 F = fun ( x )
5 fnorm = jnp . linalg . norm ( F )
6 print ( f " Iter { i }: || F || = { fnorm :.2 e } " )
7 if fnorm < kwargs . get ( ’ tol ’ , 1e -8) :
8 break
9 # ... rest of Newton step
10 return x , i

9 Advanced Topics
9.1 Handling Singular Jacobians
When the Jacobian becomes singular or ill-conditioned:

1. Pseudo-inverse approach (already implemented):

1 p = - jnp . linalg . pinv ( J ) @ F

2. Levenberg-Marquardt regularization:
1 def l e v e n b e r g_ m a rq u a rd t _ st e p (J , F , lambda_reg =1 e -6) :
2 n = J . shape [0]
3 JtJ = J . T @ J
4 JtF = J . T @ F
5 return - jnp . linalg . solve ( JtJ + lambda_reg * jnp . eye ( n ) , JtF )
 JAX Nonlinear System Solvers - Reference Manual 12

3. SVD-based solution:
1 def robust_solve (J , F , tol =1 e -10) :
2 U , S , Vt = jnp . linalg . svd (J , full_matrices = False )
3 S_inv = jnp . where ( S > tol , 1/ S , 0)
4 return - Vt . T @ ( S_inv * ( U . T @ F ) )

9.2 Continuation Methods

For difficult problems, use parameter continuation:
1 def co n ti n u at i on_solver ( fun_param , x0 , param_start , param_end , n_steps =10) :
2 """
3 Solve F ( x ; lambda ) = 0 by continuing from lambda = param_start to lambda = param_end
4 """
5 params = jnp . linspace ( param_start , param_end , n_steps )
6 x = x0
7 solutions = []
8
9 for param in params :
10 fun = lambda x : fun_param (x , param )
11 x , _ = newton ( fun , x )
12 solutions . append ( x )
13
14 return x , jnp . array ( solutions )

9.3 Parallel Solution of Multiple Systems

1 # Solve F ( x ; p ) = 0 for multiple parameters p

2 def batch_solver ( fun_param , x0 , params ) :
3 def solve_single ( p ) :
4 fun = lambda x : fun_param (x , p )
5 return newton ( fun , x0 ) [0]
6
7 # Vectorize over parameter dimension
8 vmap_solve = jax . vmap ( solve_single )
9 return vmap_solve ( params )

10 Theoretical Foundations
10.1 Convergence Analysis
Theorem 10.1 (Newton’s Method Convergence). Let F : D ⊂ Rn → Rn be continuously differentiable in
an open convex set D. Assume:

1. ∃x∗ ∈ D such that F (x∗ ) = 0

2. J(x∗ ) is nonsingular

3. J is Lipschitz continuous: ∥J(x) − J(y)∥ ≤ L ∥x − y∥

Then ∃δ > 0 such that for x0 with ∥x0 − x∗ ∥ < δ, Newton’s method converges quadratically:

L
∥xk+1 − x∗ ∥ ≤ ∥xk − x∗ ∥2 (8)
2 ∥J(x∗ )−1 ∥

Theorem 10.2 (Broyden’s Method Convergence). Under similar conditions, Broyden’s method achieves
superlinear convergence:
∥xk+1 − x∗ ∥
lim =0 (9)
k→∞ ∥xk − x∗ ∥
JAX Nonlinear System Solvers - Reference Manual 13

10.2 Trust Region Theory

The trust region subproblem:
1
minimizep m(p) = f + g T p + pT Bp subject to ∥p∥ ≤ ∆ (10)
2
Has the unique solution:
p∗ = −(B + λI)−1 g (11)
where λ ≥ 0 satisfies ∥p∗ ∥ = ∆ (if constraint is active).

10.3 Line Search Theory

10.3.1 Armijo Condition (Sufficient Decrease)

f (xk + αk pk ) ≤ f (xk ) + c1 αk ∇f (xk )T pk (12)

10.3.2 Wolfe Conditions

1. Armijo condition (above)
2. Curvature condition:
∇f (xk + αk pk )T pk ≥ c2 ∇f (xk )T pk (13)

10.3.3 Strong Wolfe Conditions

Replace curvature with:
|∇f (xk + αk pk )T pk | ≤ c2 |∇f (xk )T pk | (14)
Typical values: c1 = 10−4 , c2 = 0.9

10.4 Complexity Analysis

Table 2: Computational Complexity of Methods

Method Per-iteration Cost Storage Convergence Rate

Newton O(n3 ) O(n2 ) Quadratic
Broyden O(n2 ) O(n2 ) Superlinear
Trust Region O(n3 ) O(n2 ) Quadratic

For sparse systems, costs can be reduced significantly using sparse linear algebra.

11 References and Further Reading

11.1 Core References
1. Nocedal, J. and Wright, S.J. (2006). Numerical Optimization. 2nd Edition. Springer.
• Comprehensive treatment of line search and trust region methods
• Chapters 3, 4, and 6 are particularly relevant
2. Dennis, J.E. and Schnabel, R.B. (1996). Numerical Methods for Unconstrained Optimization and
Nonlinear Equations. SIAM.
• Classic reference for Newton and quasi-Newton methods
• Excellent coverage of globalization strategies
3. Kelley, C.T. (2003). Solving Nonlinear Equations with Newton’s Method. SIAM.
• Focused treatment of Newton-type methods
• Good practical implementation advice
 JAX Nonlinear System Solvers - Reference Manual 14

11.2 Newton-Raphson Method

4. Ortega, J.M. and Rheinboldt, W.C. (2000). Iterative Solution of Nonlinear Equations in Several
Variables. SIAM.

• Theoretical foundations of Newton’s method

• Convergence analysis and variants

5. Deuflhard, P. (2004). Newton Methods for Nonlinear Problems. Springer.

• Modern treatment with affine invariance

• Adaptive strategies and error estimation

11.3 Trust Region Methods

6. Conn, A.R., Gould, N.I.M., and Toint, P.L. (2000). Trust-Region Methods. SIAM.

• Comprehensive monograph on trust region algorithms

• Dog-leg and 2D subspace methods

7. Powell, M.J.D. (1970). “A hybrid method for nonlinear equations.” In Numerical Methods for Nonlinear
Algebraic Equations, P. Rabinowitz (ed.), Gordon and Breach.

• Original dog-leg method paper

• Combination of steepest descent and Newton

11.4 Broyden’s Method and Quasi-Newton Methods

8. Broyden, C.G. (1965). “A class of methods for solving nonlinear simultaneous equations.” Mathematics
of Computation, 19(92), 577-593.

• Original Broyden’s method paper

• Foundation of quasi-Newton methods

9. Martinez, J.M. (2000). “Practical quasi-Newton methods for solving nonlinear systems.” Journal of
Computational and Applied Mathematics, 124(1-2), 97-121.

• Modern perspective on quasi-Newton methods

• Implementation strategies and variants

10. Griewank, A. and Walther, A. (2008). Evaluating Derivatives: Principles and Techniques of Algo-
rithmic Differentiation. 2nd Edition. SIAM.

• Automatic differentiation techniques

• Essential for efficient Jacobian computation

11.5 Strong Wolfe Line Search

11. Wolfe, P. (1969). “Convergence conditions for ascent methods.” SIAM Review, 11(2), 226-235.

• Original Wolfe conditions paper

• Theoretical foundations of line search

12. Moré, J.J. and Thuente, D.J. (1994). “Line search algorithms with guaranteed sufficient decrease.”
ACM Transactions on Mathematical Software, 20(3), 286-307.

• Practical implementation of Strong Wolfe line search

• Numerical safeguards and efficiency
 JAX Nonlinear System Solvers - Reference Manual 15

11.6 JAX and Automatic Differentiation

13. Bradbury, J., et al. (2018). “JAX: Composable transformations of Python+NumPy programs.” GitHub
repository.

• JAX documentation and examples

• https://github.com/google/jax

14. Baydin, A.G., Pearlmutter, B.A., Radul, A.A., and Siskind, J.M. (2017). “Automatic differen-
tiation in machine learning: a survey.” Journal of Machine Learning Research, 18(1), 5595-5637.

• Comprehensive survey of AD techniques

• Comparison of forward and reverse mode

11.7 Numerical Analysis and Implementation

15. Higham, N.J. (2002). Accuracy and Stability of Numerical Algorithms. 2nd Edition. SIAM.

• Numerical stability considerations

• Condition number analysis

16. Press, W.H., Teukolsky, S.A., Vetterling, W.T., and Flannery, B.P. (2007). Numerical Recipes:
The Art of Scientific Computing. 3rd Edition. Cambridge University Press.

• Practical implementation advice

• Code examples and pitfalls

11.8 Specialized Topics

17. Allgower, E.L. and Georg, K. (2003). Introduction to Numerical Continuation Methods. SIAM.

• Continuation and homotopy methods

• Path following algorithms

18. Facchinei, F. and Pang, J.S. (2003). Finite-Dimensional Variational Inequalities and Complementarity
Problems. Springer.

• Extensions to constrained problems

• Semismooth Newton methods

11.9 Recent Advances

19. Cartis, C., Gould, N.I.M., and Toint, P.L. (2011). “Adaptive cubic regularisation methods for
unconstrained optimization.” Mathematical Programming, 127(2), 245-295.

• Modern globalization strategies

• Worst-case complexity bounds

20. Bergou, E., Diouane, Y., and Kungurtsev, V. (2020). “Convergence and complexity of Newton
iteration for operator equations.” Journal of Computational and Applied Mathematics, 372, 112723.

• Recent theoretical advances

• Complexity analysis for Newton methods
 JAX Nonlinear System Solvers - Reference Manual 16

11.10 Online Resources

21. NEOS Guide to Optimization

• https://neos-guide.org/
• Interactive examples and solver comparisons

22. Matrix Computation Toolbox

• https://www.mathworks.com/matlabcentral/fileexchange/1322
• MATLAB implementations of various algorithms

23. SciPy Optimization Tutorial

• https://docs.scipy.org/doc/scipy/tutorial/optimize.html
• Python implementations and examples

11.11 Software Libraries

24. MINPACK - Classic Fortran library for nonlinear equations

• Powell hybrid method implementation

• Levenberg-Marquardt algorithm

25. PETSc - Portable, Extensible Toolkit for Scientific Computation

• Scalable nonlinear solvers

• MPI parallel implementation

26. SUNDIALS - SUite of Nonlinear and DIfferential/ALgebraic equation Solvers

• KINSOL for nonlinear systems

• Advanced preconditioning options

A Complete Code Example

Here’s a complete example demonstrating advanced usage:
1 import jax
2 import jax . numpy as jnp
3 import matplotlib . pyplot as plt
4 from nonlin ear_solvers import newton , broyden_hybrid
5 import time
6
7 # Enable 64 - bit precision
8 jax . config . update ( " jax_enable_x64 " , True )
9
10 # Define a challenging test problem : Chandrasekhar H - equation
11 def chandrasekhar_H (x , c =0.9) :
12 """
13 Discretized Chandrasekhar H - equation :
14 H ( mu ) = 1 + ( c /2) H ( mu ) int_0 ^1 [ mu /( mu + nu ) ] H ( nu ) dnu
15 """
16 n = len ( x )
17 mu = jnp . linspace (0 , 1 , n , endpoint = False ) + 0.5/ n
18 F = jnp . zeros ( n )
19
20 for i in range ( n ) :
21 integral = 0.0
22 for j in range ( n ) :
23 integral += mu [ i ] / ( mu [ i ] + mu [ j ]) * x [ j ] / n
24 F = F . at [ i ]. set ( x [ i ] - 1 - 0.5 * c * x [ i ] * integral )
25
26 return F
 JAX Nonlinear System Solvers - Reference Manual 17

27
28 # Problem sizes to test
29 sizes = [10 , 20 , 50 , 100]
30 results = {}
31
32 for n in sizes :
33 print ( f " \ nSolving Chandrasekhar H - equation with n ={ n } " )
34
35 # Initial guess
36 x0 = jnp . ones ( n )
37
38 # Newton with line search
39 t0 = time . time ()
40 x_newton_ls , iter_newton_ls = newton (
41 lambda x : chandrasekhar_H (x , c =0.9) ,
42 x0 ,
43 method = ’ line_search ’
44 )
45 time_newton_ls = time . time () - t0
46
47 # Newton with trust region
48 t0 = time . time ()
49 x_newton_tr , iter_newton_tr = newton (
50 lambda x : chandrasekhar_H (x , c =0.9) ,
51 x0 ,
52 method = ’ trust_region ’
53 )
54 time_newton_tr = time . time () - t0
55
56 # Broyden
57 t0 = time . time ()
58 x_broyden , iter_broyden = broyden_hybrid (
59 lambda x : chandrasekhar_H (x , c =0.9) ,
60 x0 ,
61 refresh_freq = max (5 , n //10)
62 )
63 time_broyden = time . time () - t0
64
65 # Store results
66 results [ n ] = {
67 ’ newton_ls ’: ( iter_newton_ls , time_newton_ls , x_newton_ls ) ,
68 ’ newton_tr ’: ( iter_newton_tr , time_newton_tr , x_newton_tr ) ,
69 ’ broyden ’: ( iter_broyden , time_broyden , x_broyden )
70 }
71
72 # Print summary
73 print ( f " Newton ( LS ) : { iter_newton_ls } iters , { time_newton_ls :.3 f } sec " )
74 print ( f " Newton ( TR ) : { iter_newton_tr } iters , { time_newton_tr :.3 f } sec " )
75 print ( f " Broyden : { iter_broyden } iters , { time_broyden :.3 f } sec " )
76
77 # Verify solution quality
78 n = 50
79 x_sol = results [ n ][ ’ newton_ls ’ ][2]
80 residual = chandrasekhar_H ( x_sol , c =0.9)
81 print ( f " \ nSolution quality check ( n ={ n }) : " )
82 print ( f " || F ( x ) || _inf = { jnp . max ( jnp . abs ( residual ) ) :.2 e } " )
83 print ( f " || F ( x ) || _2 = { jnp . linalg . norm ( residual ) :.2 e } " )
Listing 1: Chandrasekhar H-equation solver

B Common Patterns and Best Practices

B.1 Problem Scaling

1 def solve_scaled ( fun , x0 , x_scale = None , f_scale = None , ** kwargs ) :

2 """
3 Solve with automatic scaling for better conditioning
 JAX Nonlinear System Solvers - Reference Manual 18

4 """
5 if x_scale is None :
6 x_scale = jnp . maximum ( jnp . abs ( x0 ) , 1.0)
7 if f_scale is None :
8 f_scale = jnp . maximum ( jnp . abs ( fun ( x0 ) ) , 1.0)
9
10 # Scaled function
11 def fun_scaled ( x_scaled ) :
12 x = x_scaled * x_scale
13 return fun ( x ) / f_scale
14
15 # Solve scaled problem
16 x_scaled_sol , iters = newton ( fun_scaled , x0 / x_scale , ** kwargs )
17
18 # Unscale solution
19 return x_scaled_sol * x_scale , iters
Listing 2: Automatic scaling wrapper

B.2 Robustness Wrapper

1 def robust_solve ( fun , x0 , methods =[ ’ newton_ls ’ , ’ newton_tr ’ , ’ broyden ’] , ** kwargs ) :

2 """
3 Try multiple methods and return the best result
4 """
5 best_x = None
6 best_residual = jnp . inf
7
8 for method in methods :
9 try :
10 if method == ’ newton_ls ’:
11 x , _ = newton ( fun , x0 , method = ’ line_search ’ , ** kwargs )
12 elif method == ’ newton_tr ’:
13 x , _ = newton ( fun , x0 , method = ’ trust_region ’ , ** kwargs )
14 elif method == ’ broyden ’:
15 x , _ = broyden_hybrid ( fun , x0 , ** kwargs )
16
17 residual = jnp . linalg . norm ( fun ( x ) )
18 if residual < best_residual :
19 best_residual = residual
20 best_x = x
21
22 except Exception as e :
23 print ( f " Method { method } failed : { e } " )
24 continue
25
26 return best_x , best_residual
Listing 3: Multi-method robust solver

B.3 Convergence Monitoring

1 class Con ve rg enceMonitor :

2 def __init__ ( self ) :
3 self . history = []
4
5 def __call__ ( self , x , f_norm , iteration ) :
6 self . history . append ({
7 ’ iteration ’: iteration ,
8 ’x ’: x . copy () ,
9 ’ f_norm ’: f_norm ,
10 ’ time ’: time . time ()
11 })
12
13 def plot ( self ) :
14 iterations = [ h [ ’ iteration ’] for h in self . history ]
 JAX Nonlinear System Solvers - Reference Manual 19

15 f_norms = [ h [ ’ f_norm ’] for h in self . history ]

16
17 plt . figure ( figsize =(10 , 6) )
18 plt . semilogy ( iterations , f_norms , ’b - o ’)
19 plt . xlabel ( ’ Iteration ’)
20 plt . ylabel ( ’ || F ( x ) || ’)
21 plt . title ( ’ Convergence History ’)
22 plt . grid ( True )
23 plt . show ()
Listing 4: Convergence monitoring class

C Troubleshooting Checklist
When encountering issues, work through this checklist:

1. Verify the problem formulation

□ Is F (x) correctly implemented?

□ Are there typos in the equations?
□ Is the expected solution known?

2. Check the initial guess

□ Is x0 in the domain of F ?
□ Is x0 reasonably close to a solution?
□ Try multiple random initial guesses

3. Examine problem conditioning

□ Compute condition number of Jacobian

□ Check for near-singular points
□ Consider scaling or reformulation

4. Adjust solver parameters

□ Try different values of tol

□ Increase maxiter
□ Modify line search parameters

5. Switch methods

□ If Newton fails, try trust region

□ If Jacobian expensive, try Broyden
□ Consider hybrid approaches

6. Debug the solver

□ Enable JAX debugging

□ Add convergence monitoring
□ Check intermediate values
JAX Nonlinear System Solvers - Reference Manual 20

This comprehensive manual provides everything needed to effectively use and understand the
JAX-based nonlinear system solvers. The combination of theoretical background, practical
examples, and extensive references makes it suitable both as a user guide and as a reference for
further study.