# Path: diffrax/custom_types.py
# def _item_to_str(item: Union[str, type, slice]) -> str:
# def _maybe_tuple_to_str(
#     item: Union[str, type, slice, Tuple[Union[str, type, slice], ...]]
# ) -> str:
#     def __class_getitem__(cls, item):
#     def __class_getitem__(cls, item):
#     def __class_getitem__(cls, item):
#     def __class_getitem__(cls, item):
# class _Array(Generic[_Annotation]):
# class _PyTree(Generic[_Annotation]):
# class Array:
#         class X:
# class PyTree:
#         class X:
# class Scalar:
# class Array:
# class PyTree:
# 
# Path: diffrax/misc/errors.py
# def error_if(
#     pred: Union[_Bool, Callable[[], _Bool]],
#     msg: str,
#     error_cls: Type[Exception] = ValueError,
# ) -> bool:
#     """For use as part of validating inputs.
#     Works even under JIT.
# 
#     Example:
#         @jax.jit
#         def f(x):
#             error_if(x < 0, "x must be >= 0")
# 
#         f(jax.numpy.array(-1))
#     """
#     branched_error_if(pred, 0, [msg], error_cls)
# 
# Path: diffrax/misc/misc.py
# def rms_norm(x: PyTree) -> Scalar:
#     x, _ = fu.ravel_pytree(x)
#     if x.size == 0:
#         return 0
#     sqnorm = jnp.mean(x**2)
#     cond = sqnorm == 0
#     # Double-where trick to avoid NaN gradients.
#     # See JAX issues #5039 and #1052.
#     _sqnorm = jnp.where(cond, 1.0, sqnorm)
#     return jnp.where(cond, 0.0, jnp.sqrt(_sqnorm))
# 
# Path: diffrax/solution.py
# class RESULTS(metaclass=ContainerMeta):
#     successful = ""
#     max_steps_reached = (
#         "The maximum number of solver steps was reached. Try increasing `max_steps`."
#     )
#     dt_min_reached = "The minimum step size was reached."
#     nan_time = "NaN time encountered during timestepping."
#     implicit_divergence = "Implicit method diverged."
#     implicit_nonconvergence = (
#         "Implicit method did not converge within the required number of iterations."
#     )
# 
# Path: diffrax/nonlinear_solver/base.py
# class NonlinearSolution(eqx.Module):
# class AbstractNonlinearSolver(eqx.Module):
#     def __init_subclass__(cls, **kwargs):
#     def _solve(
#         self,
#         fn: Callable,
#         x: PyTree,
#         jac: Optional[LU_Jacobian],
#         nondiff_args: PyTree,
#         diff_args: PyTree,
#     ) -> Tuple[PyTree, RESULTS]:
#     def __call__(
#         self, fn: Callable, x: PyTree, args: PyTree, jac: Optional[LU_Jacobian] = None
#     ) -> NonlinearSolution:
#     def jac(fn: Callable, x: PyTree, args: PyTree) -> LU_Jacobian:
# def _root_solve_jvp(
#     self: AbstractNonlinearSolver,
#     fn: callable,
#     x: PyTree,
#     jac: Optional[LU_Jacobian],
#     nondiff_args: PyTree,
#     diff_args: PyTree,
#     tang_diff_args: PyTree,
# ):
#     def _for_jac(_root):
#     def _for_jvp(_diff_args):


from typing import Callable, Optional, Tuple
from ..custom_types import Bool, Int, PyTree, Scalar
from ..misc import error_if, rms_norm
from ..solution import RESULTS
from .base import AbstractNonlinearSolver, LU_Jacobian, NonlinearSolution
import equinox as eqx
import jax
import jax.flatten_util as fu
import jax.lax as lax
import jax.numpy as jnp
import jax.scipy as jsp
            val = (flat, step + 1, diffsize, diffsize_prev)
            return val

        # Unconditionally execute two loops to fill in diffsize and diffsize_prev.
        val = (flat, 0, None, None)
        val = body_fn(val)
        val = body_fn(val)
        val = lax.while_loop(cond_fn, body_fn, val)
        flat, num_steps, diffsize, diffsize_prev = val

        if self.tolerate_nonconvergence:
            result = RESULTS.successful
        else:
            rate = diffsize / diffsize_prev
            factor = diffsize * rate / (1 - rate)
            converged = _converged(factor, self.kappa)
            diverged = _diverged(rate)
            small = _small(diffsize)
            result = jnp.where(
                converged, RESULTS.successful, RESULTS.implicit_nonconvergence
            )
            result = jnp.where(diverged, RESULTS.implicit_divergence, result)
            result = jnp.where(small, RESULTS.successful, result)
        root = unflatten(flat)
        return NonlinearSolution(root=root, num_steps=num_steps, result=result)


NewtonNonlinearSolver.__init__.__doc__ = """
**Arguments:**


# Path: diffrax/custom_types.py
# def _item_to_str(item: Union[str, type, slice]) -> str:
# def _maybe_tuple_to_str(
#     item: Union[str, type, slice, Tuple[Union[str, type, slice], ...]]
# ) -> str:
#     def __class_getitem__(cls, item):
#     def __class_getitem__(cls, item):
#     def __class_getitem__(cls, item):
#     def __class_getitem__(cls, item):
# class _Array(Generic[_Annotation]):
# class _PyTree(Generic[_Annotation]):
# class Array:
#         class X:
# class PyTree:
#         class X:
# class Scalar:
# class Array:
# class PyTree:
# 
# Path: diffrax/misc/errors.py
# def error_if(
#     pred: Union[_Bool, Callable[[], _Bool]],
#     msg: str,
#     error_cls: Type[Exception] = ValueError,
# ) -> bool:
#     """For use as part of validating inputs.
#     Works even under JIT.
# 
#     Example:
#         @jax.jit
#         def f(x):
#             error_if(x < 0, "x must be >= 0")
# 
#         f(jax.numpy.array(-1))
#     """
#     branched_error_if(pred, 0, [msg], error_cls)
# 
# Path: diffrax/misc/misc.py
# def rms_norm(x: PyTree) -> Scalar:
#     x, _ = fu.ravel_pytree(x)
#     if x.size == 0:
#         return 0
#     sqnorm = jnp.mean(x**2)
#     cond = sqnorm == 0
#     # Double-where trick to avoid NaN gradients.
#     # See JAX issues #5039 and #1052.
#     _sqnorm = jnp.where(cond, 1.0, sqnorm)
#     return jnp.where(cond, 0.0, jnp.sqrt(_sqnorm))
# 
# Path: diffrax/solution.py
# class RESULTS(metaclass=ContainerMeta):
#     successful = ""
#     max_steps_reached = (
#         "The maximum number of solver steps was reached. Try increasing `max_steps`."
#     )
#     dt_min_reached = "The minimum step size was reached."
#     nan_time = "NaN time encountered during timestepping."
#     implicit_divergence = "Implicit method diverged."
#     implicit_nonconvergence = (
#         "Implicit method did not converge within the required number of iterations."
#     )
# 
# Path: diffrax/nonlinear_solver/base.py
# class NonlinearSolution(eqx.Module):
# class AbstractNonlinearSolver(eqx.Module):
#     def __init_subclass__(cls, **kwargs):
#     def _solve(
#         self,
#         fn: Callable,
#         x: PyTree,
#         jac: Optional[LU_Jacobian],
#         nondiff_args: PyTree,
#         diff_args: PyTree,
#     ) -> Tuple[PyTree, RESULTS]:
#     def __call__(
#         self, fn: Callable, x: PyTree, args: PyTree, jac: Optional[LU_Jacobian] = None
#     ) -> NonlinearSolution:
#     def jac(fn: Callable, x: PyTree, args: PyTree) -> LU_Jacobian:
# def _root_solve_jvp(
#     self: AbstractNonlinearSolver,
#     fn: callable,
#     x: PyTree,
#     jac: Optional[LU_Jacobian],
#     nondiff_args: PyTree,
#     diff_args: PyTree,
#     tang_diff_args: PyTree,
# ):
#     def _for_jac(_root):
#     def _for_jvp(_diff_args):


# Path: diffrax/nonlinear_solver/newton.py
from typing import Callable, Optional, Tuple
from ..custom_types import Bool, Int, PyTree, Scalar
from ..misc import error_if, rms_norm
from ..solution import RESULTS
from .base import AbstractNonlinearSolver, LU_Jacobian, NonlinearSolution
import equinox as eqx
import jax
import jax.flatten_util as fu
import jax.lax as lax
import jax.numpy as jnp
import jax.scipy as jsp

            val = (flat, step + 1, diffsize, diffsize_prev)
            return val

        # Unconditionally execute two loops to fill in diffsize and diffsize_prev.
        val = (flat, 0, None, None)
        val = body_fn(val)
        val = body_fn(val)
        val = lax.while_loop(cond_fn, body_fn, val)
        flat, num_steps, diffsize, diffsize_prev = val

        if self.tolerate_nonconvergence:
            result = RESULTS.successful
        else:
            rate = diffsize / diffsize_prev
            factor = diffsize * rate / (1 - rate)
            converged = _converged(factor, self.kappa)
            diverged = _diverged(rate)
            small = _small(diffsize)
            result = jnp.where(
                converged, RESULTS.successful, RESULTS.implicit_nonconvergence
            )
            result = jnp.where(diverged, RESULTS.implicit_divergence, result)
            result = jnp.where(small, RESULTS.successful, result)
        root = unflatten(flat)
        return NonlinearSolution(root=root, num_steps=num_steps, result=result)


NewtonNonlinearSolver.__init__.__doc__ = """
**Arguments:**


- `max_steps`: The maximum number of steps allowed. If more than this are required then