Allow for dynamic `atol`

I am trying to use a Newton solver inside a diffrax simulation to solve a problem that is nearly linear (and is sometimes completely linear but I can't guarantee this in advance, think of a linearly interpolated tabular function) and would like to use cauchy_termination. However, I don't know the scale of my values in advance and this may change from cell to cell or over the course of the simulation. I have tried to address this by modifying the objective function but the increased nonlinearity causes issues.

As an example (not really an MWE as it doesn't really demonstrate the problem), let's use an actually linear objective function

solver = optx.Newton(atol=1e-14, rtol=0.0, linear_solver=lx.Diagonal())

def linear_objective_function(y, args):
    return y - args["y_target"]

def relative_objective_function(y, args):
    return y / args["y_target"] - 1.0

sol = optx.root_find(
    ***_objective_function,
    solver,
    jnp.logspace(0.0, 12.0, num=100),
    args={"y_target": 0.9 * jnp.logspace(0.0, 12.0, num=100)},
    tags=frozenset({lx.diagonal_tag}),
)  

The linear objective function theoretically only requires one step (although sol.state.step = 2) to get to the optimal solution but due to floating point error and a single fixed atol may not converge in a real-world case due to to tolerance being set to low or be flagged as having converged too early in a nonlinear setting due to tolerance being set to high.

However, the relative function requires 4 steps (or more in a slightly nonlinear case) and in my testing shows larger discrepancies against finite difference even with fairly large relative bumps in y_target.

atol, rtol are type-hinted as float not array-like so my assumption is that they cannot be tracers.

I tried cauchy_termination=False where I can use rtol but this requires at least two iterations (unnecessary when the problem is actually linear) and caused another issue I can't remember off the top of my head (although this could have been coming from elsewhere and something I've since fixed).

Start by reading the Newton solver and its cauchy_termination handling, then reproduce the provided root_find example with the linear and relative objective functions. Determine how dynamic atol could follow changing value scales without introducing the relative formulation's discrepancies, and verify that the solver converges appropriately for both linear and slightly nonlinear cases.