# 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/ad.py
# class fixed_custom_jvp:
#     """As jax.custom_jvp but works around JAX issue #9374."""
# 
#     def __init__(self, fn, nondiff_argnums=()):
#         assert set(nondiff_argnums) == set(range(len(nondiff_argnums)))
# 
#         def fn_wrapper(nondiff_args_nontracer, nondiff_args_tracer, diff_args):
#             nondiff_args = eqx.combine(nondiff_args_nontracer, nondiff_args_tracer)
#             return fn(*nondiff_args, *diff_args)
# 
#         self.fn = jax.custom_jvp(fn_wrapper, nondiff_argnums=(0,))
#         self.cutoff = max(nondiff_argnums, default=-1) + 1
#         self.fn_jvp = None
# 
#     def defjvp(self, fn_jvp):
#         def fn_jvp_wrapper(nondiff_args_nontracer, combined_args, tang_combined_args):
#             nondiff_args_tracer, diff_args = combined_args
#             _, tang_diff_args = tang_combined_args
#             nondiff_args = eqx.combine(nondiff_args_nontracer, nondiff_args_tracer)
#             return fn_jvp(*nondiff_args, diff_args, tang_diff_args)
# 
#         self.fn.defjvp(fn_jvp_wrapper)
# 
#     def __call__(self, *args):
#         is_tracer = lambda x: isinstance(x, jax.core.Tracer)
#         nondiff_args = args[: self.cutoff]
#         diff_args = args[self.cutoff :]
#         nondiff_args_tracer, nondiff_args_nontracer = eqx.partition(
#             nondiff_args, is_tracer
#         )
#         nondiff_args_tracer = jax.tree_map(lax.stop_gradient, nondiff_args_tracer)
#         return self.fn(nondiff_args_nontracer, nondiff_args_tracer, diff_args)
# 
# 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."
#     )


import abc
import equinox as eqx
import jax
import jax.flatten_util as fu
import jax.numpy as jnp
import jax.scipy as jsp
from typing import Callable, Optional, Tuple, TypeVar
from ..custom_types import Int, PyTree, Scalar
from ..misc import fixed_custom_jvp
from ..solution import RESULTS
class NonlinearSolution(eqx.Module):
    root: PyTree
    num_steps: Int
    result: RESULTS


class AbstractNonlinearSolver(eqx.Module):
    """Abstract base class for all nonlinear root-finding algorithms.

    Subclasses will be differentiable via the implicit function theorem.
    """

    rtol: Optional[Scalar] = None
    atol: Optional[Scalar] = None

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # Note that this breaks the descriptor protocol so we have to pass self
        # manually in __call__.
        cls._solve = fixed_custom_jvp(cls._solve, nondiff_argnums=(0, 1, 2, 3, 4))
        cls._solve.defjvp(_root_solve_jvp)

    @abc.abstractmethod
    def _solve(
        self,
        fn: Callable,
        x: PyTree,
        jac: Optional[LU_Jacobian],
        nondiff_args: PyTree,
        diff_args: PyTree,

# 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/ad.py
# class fixed_custom_jvp:
#     """As jax.custom_jvp but works around JAX issue #9374."""
# 
#     def __init__(self, fn, nondiff_argnums=()):
#         assert set(nondiff_argnums) == set(range(len(nondiff_argnums)))
# 
#         def fn_wrapper(nondiff_args_nontracer, nondiff_args_tracer, diff_args):
#             nondiff_args = eqx.combine(nondiff_args_nontracer, nondiff_args_tracer)
#             return fn(*nondiff_args, *diff_args)
# 
#         self.fn = jax.custom_jvp(fn_wrapper, nondiff_argnums=(0,))
#         self.cutoff = max(nondiff_argnums, default=-1) + 1
#         self.fn_jvp = None
# 
#     def defjvp(self, fn_jvp):
#         def fn_jvp_wrapper(nondiff_args_nontracer, combined_args, tang_combined_args):
#             nondiff_args_tracer, diff_args = combined_args
#             _, tang_diff_args = tang_combined_args
#             nondiff_args = eqx.combine(nondiff_args_nontracer, nondiff_args_tracer)
#             return fn_jvp(*nondiff_args, diff_args, tang_diff_args)
# 
#         self.fn.defjvp(fn_jvp_wrapper)
# 
#     def __call__(self, *args):
#         is_tracer = lambda x: isinstance(x, jax.core.Tracer)
#         nondiff_args = args[: self.cutoff]
#         diff_args = args[self.cutoff :]
#         nondiff_args_tracer, nondiff_args_nontracer = eqx.partition(
#             nondiff_args, is_tracer
#         )
#         nondiff_args_tracer = jax.tree_map(lax.stop_gradient, nondiff_args_tracer)
#         return self.fn(nondiff_args_nontracer, nondiff_args_tracer, diff_args)
# 
# 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
import abc
import equinox as eqx
import jax
import jax.flatten_util as fu
import jax.numpy as jnp
import jax.scipy as jsp
from typing import Callable, Optional, Tuple, TypeVar
from ..custom_types import Int, PyTree, Scalar
from ..misc import fixed_custom_jvp
from ..solution import RESULTS

class NonlinearSolution(eqx.Module):
    root: PyTree
    num_steps: Int
    result: RESULTS


class AbstractNonlinearSolver(eqx.Module):
    """Abstract base class for all nonlinear root-finding algorithms.

    Subclasses will be differentiable via the implicit function theorem.
    """

    rtol: Optional[Scalar] = None
    atol: Optional[Scalar] = None

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # Note that this breaks the descriptor protocol so we have to pass self
        # manually in __call__.
        cls._solve = fixed_custom_jvp(cls._solve, nondiff_argnums=(0, 1, 2, 3, 4))
        cls._solve.defjvp(_root_solve_jvp)

    @abc.abstractmethod
    def _solve(
        self,
        fn: Callable,
        x: PyTree,
        jac: Optional[LU_Jacobian],
        nondiff_args: PyTree,
        diff_args: PyTree,

    ) -> Tuple[PyTree, RESULTS]: