# 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
    !!! warning

        Note that backpropagation through `__call__` may not produce accurate values if
        `tolerate_nonconvergence=True`, as the backpropagation calculation implicitly
        assumes that the forward pass converged.
    """

    max_steps: Optional[Int] = 10
    kappa: Scalar = 1e-2
    norm: Callable = rms_norm
    tolerate_nonconvergence: bool = False

    def __post_init__(self):
        if self.max_steps is not None:
            error_if(self.max_steps < 2, "max_steps must be at least 2.")

    def _solve(
        self,
        fn: callable,
        x: PyTree,
        jac: Optional[LU_Jacobian],
        nondiff_args: PyTree,
        diff_args: PyTree,
    ) -> Tuple[PyTree, RESULTS]:
        args = eqx.combine(nondiff_args, diff_args)
        rtol = 1e-3 if self.rtol is None else self.rtol
        atol = 1e-6 if self.atol is None else self.atol
        scale = atol + rtol * self.norm(x)
        flat, unflatten = fu.ravel_pytree(x)
        if flat.size == 0:

# 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

    !!! warning

        Note that backpropagation through `__call__` may not produce accurate values if
        `tolerate_nonconvergence=True`, as the backpropagation calculation implicitly
        assumes that the forward pass converged.
    """

    max_steps: Optional[Int] = 10
    kappa: Scalar = 1e-2
    norm: Callable = rms_norm
    tolerate_nonconvergence: bool = False

    def __post_init__(self):
        if self.max_steps is not None:
            error_if(self.max_steps < 2, "max_steps must be at least 2.")

    def _solve(
        self,
        fn: callable,
        x: PyTree,
        jac: Optional[LU_Jacobian],
        nondiff_args: PyTree,
        diff_args: PyTree,
    ) -> Tuple[PyTree, RESULTS]:
        args = eqx.combine(nondiff_args, diff_args)
        rtol = 1e-3 if self.rtol is None else self.rtol
        atol = 1e-6 if self.atol is None else self.atol
        scale = atol + rtol * self.norm(x)
        flat, unflatten = fu.ravel_pytree(x)
        if flat.size == 0:

            return NonlinearSolution(root=x, num_steps=0, result=RESULTS.successful)