# Path: diffrax/local_interpolation.py
# class ThirdOrderHermitePolynomialInterpolation(AbstractLocalInterpolation):
#     coeffs: PyTree[Array[4, "dims":...]]  # noqa: F821
# 
#     def __init__(
#         self,
#         *,
#         y0: PyTree[Array["dims":...]],  # noqa: F821
#         y1: PyTree[Array["dims":...]],  # noqa: F821
#         f0: PyTree[Array["dims":...]],  # noqa: F821
#         f1: PyTree[Array["dims":...]],  # noqa: F821
#         **kwargs
#     ):
#         super().__init__(**kwargs)
# 
#         def _calculate(_y0, _y1, _f0, _f1):
#             _a = _f0 + _f1 + 2 * _y0 - 2 * _y1
#             _b = -2 * _f0 - _f1 - 3 * _y0 + 3 * _y1
#             return jnp.stack([_a, _b, _f0, _y0])
# 
#         self.coeffs = jax.tree_map(_calculate, y0, y1, f0, f1)
# 
#     @classmethod
#     def from_k(
#         cls,
#         *,
#         y0: PyTree[Array["dims":...]],  # noqa: F821
#         y1: PyTree[Array["dims":...]],  # noqa: F821
#         k: PyTree[Array["order", "dims":...]],  # noqa: F821
#         **kwargs
#     ):
#         return cls(y0=y0, y1=y1, f0=ω(k)[0].ω, f1=ω(k)[-1].ω, **kwargs)
# 
#     def evaluate(
#         self, t0: Scalar, t1: Optional[Scalar] = None, left: bool = True
#     ) -> PyTree:
#         del left
#         if t1 is not None:
#             return self.evaluate(t1) - self.evaluate(t0)
# 
#         t = linear_rescale(self.t0, t0, self.t1)
# 
#         def _eval(_coeffs):
#             return jnp.polyval(_coeffs, t)
# 
#         return jax.tree_map(_eval, self.coeffs)
# 
# Path: diffrax/solver/runge_kutta.py
# class AbstractERK(AbstractRungeKutta):
#     def _recompute_jac(self, i: int) -> bool:
#         assert False
# 
# class ButcherTableau:
#     # Explicit RK methods
#     c: np.ndarray
#     b_sol: np.ndarray
#     b_error: np.ndarray
#     a_lower: Tuple[np.ndarray]
# 
#     # Implicit RK methods
#     a_diagonal: Optional[np.ndarray] = None
#     a_predictor: Optional[Tuple[np.ndarray]] = None
# 
#     # Determine the use of fast-paths
#     ssal: bool = field(init=False)
#     fsal: bool = field(init=False)
# 
#     def __post_init__(self):
#         assert self.c.ndim == 1
#         for a_i in self.a_lower:
#             assert a_i.ndim == 1
#         assert self.b_sol.ndim == 1
#         assert self.b_error.ndim == 1
#         assert self.c.shape[0] == len(self.a_lower)
#         assert all(i + 1 == a_i.shape[0] for i, a_i in enumerate(self.a_lower))
#         assert self.c.shape[0] + 1 == self.b_sol.shape[0]
#         assert self.c.shape[0] + 1 == self.b_error.shape[0]
#         for i, (a_i, c_i) in enumerate(zip(self.a_lower, self.c)):
#             diagonal = 0 if self.a_diagonal is None else self.a_diagonal[i + 1]
#             assert np.allclose(sum(a_i) + diagonal, c_i)
#         assert np.allclose(sum(self.b_sol), 1.0)
#         assert np.allclose(sum(self.b_error), 0.0)
# 
#         if self.a_diagonal is None:
#             assert self.a_predictor is None
#             implicit = False
#         else:
#             assert self.a_predictor is not None
#             implicit = True
#         if implicit:
#             assert self.a_diagonal.ndim == 1
#             assert self.c.shape[0] + 1 == self.a_diagonal.shape[0]
#             assert len(self.a_lower) == len(self.a_predictor)
#             for a_lower_i, a_predictor_i in zip(self.a_lower, self.a_predictor):
#                 assert a_lower_i.shape == a_predictor_i.shape
#                 assert np.allclose(sum(a_predictor_i), 1.0)
# 
#         lower_b_sol_equal = (self.b_sol[:-1] == self.a_lower[-1]).all()
#         last_diagonal = 0 if self.a_diagonal is None else self.a_diagonal[-1]
#         diagonal_b_sol_equal = self.b_sol[-1] == last_diagonal
#         explicit_first_stage = self.a_diagonal is None or (self.a_diagonal[0] == 0)
#         explicit_last_stage = self.a_diagonal is None or (self.a_diagonal[-1] == 0)
#         # Solution y1 is the same as the last stage
#         object.__setattr__(
#             self,
#             "ssal",
#             lower_b_sol_equal and diagonal_b_sol_equal and explicit_last_stage,
#         )
#         # Vector field - control product k1 is the same across first/last stages.
#         object.__setattr__(
#             self,
#             "fsal",
#             lower_b_sol_equal and diagonal_b_sol_equal and explicit_first_stage,
#         )


import numpy as np
from ..local_interpolation import ThirdOrderHermitePolynomialInterpolation
from .runge_kutta import AbstractERK, ButcherTableau



_fehlberg2_tableau = ButcherTableau(
    a_lower=(np.array([1 / 2]), np.array([1 / 256, 255 / 256])),
    b_sol=np.array([1 / 512, 255 / 256, 1 / 512]),

# Path: diffrax/local_interpolation.py
# class ThirdOrderHermitePolynomialInterpolation(AbstractLocalInterpolation):
#     coeffs: PyTree[Array[4, "dims":...]]  # noqa: F821
# 
#     def __init__(
#         self,
#         *,
#         y0: PyTree[Array["dims":...]],  # noqa: F821
#         y1: PyTree[Array["dims":...]],  # noqa: F821
#         f0: PyTree[Array["dims":...]],  # noqa: F821
#         f1: PyTree[Array["dims":...]],  # noqa: F821
#         **kwargs
#     ):
#         super().__init__(**kwargs)
# 
#         def _calculate(_y0, _y1, _f0, _f1):
#             _a = _f0 + _f1 + 2 * _y0 - 2 * _y1
#             _b = -2 * _f0 - _f1 - 3 * _y0 + 3 * _y1
#             return jnp.stack([_a, _b, _f0, _y0])
# 
#         self.coeffs = jax.tree_map(_calculate, y0, y1, f0, f1)
# 
#     @classmethod
#     def from_k(
#         cls,
#         *,
#         y0: PyTree[Array["dims":...]],  # noqa: F821
#         y1: PyTree[Array["dims":...]],  # noqa: F821
#         k: PyTree[Array["order", "dims":...]],  # noqa: F821
#         **kwargs
#     ):
#         return cls(y0=y0, y1=y1, f0=ω(k)[0].ω, f1=ω(k)[-1].ω, **kwargs)
# 
#     def evaluate(
#         self, t0: Scalar, t1: Optional[Scalar] = None, left: bool = True
#     ) -> PyTree:
#         del left
#         if t1 is not None:
#             return self.evaluate(t1) - self.evaluate(t0)
# 
#         t = linear_rescale(self.t0, t0, self.t1)
# 
#         def _eval(_coeffs):
#             return jnp.polyval(_coeffs, t)
# 
#         return jax.tree_map(_eval, self.coeffs)
# 
# Path: diffrax/solver/runge_kutta.py
# class AbstractERK(AbstractRungeKutta):
#     def _recompute_jac(self, i: int) -> bool:
#         assert False
# 
# class ButcherTableau:
#     # Explicit RK methods
#     c: np.ndarray
#     b_sol: np.ndarray
#     b_error: np.ndarray
#     a_lower: Tuple[np.ndarray]
# 
#     # Implicit RK methods
#     a_diagonal: Optional[np.ndarray] = None
#     a_predictor: Optional[Tuple[np.ndarray]] = None
# 
#     # Determine the use of fast-paths
#     ssal: bool = field(init=False)
#     fsal: bool = field(init=False)
# 
#     def __post_init__(self):
#         assert self.c.ndim == 1
#         for a_i in self.a_lower:
#             assert a_i.ndim == 1
#         assert self.b_sol.ndim == 1
#         assert self.b_error.ndim == 1
#         assert self.c.shape[0] == len(self.a_lower)
#         assert all(i + 1 == a_i.shape[0] for i, a_i in enumerate(self.a_lower))
#         assert self.c.shape[0] + 1 == self.b_sol.shape[0]
#         assert self.c.shape[0] + 1 == self.b_error.shape[0]
#         for i, (a_i, c_i) in enumerate(zip(self.a_lower, self.c)):
#             diagonal = 0 if self.a_diagonal is None else self.a_diagonal[i + 1]
#             assert np.allclose(sum(a_i) + diagonal, c_i)
#         assert np.allclose(sum(self.b_sol), 1.0)
#         assert np.allclose(sum(self.b_error), 0.0)
# 
#         if self.a_diagonal is None:
#             assert self.a_predictor is None
#             implicit = False
#         else:
#             assert self.a_predictor is not None
#             implicit = True
#         if implicit:
#             assert self.a_diagonal.ndim == 1
#             assert self.c.shape[0] + 1 == self.a_diagonal.shape[0]
#             assert len(self.a_lower) == len(self.a_predictor)
#             for a_lower_i, a_predictor_i in zip(self.a_lower, self.a_predictor):
#                 assert a_lower_i.shape == a_predictor_i.shape
#                 assert np.allclose(sum(a_predictor_i), 1.0)
# 
#         lower_b_sol_equal = (self.b_sol[:-1] == self.a_lower[-1]).all()
#         last_diagonal = 0 if self.a_diagonal is None else self.a_diagonal[-1]
#         diagonal_b_sol_equal = self.b_sol[-1] == last_diagonal
#         explicit_first_stage = self.a_diagonal is None or (self.a_diagonal[0] == 0)
#         explicit_last_stage = self.a_diagonal is None or (self.a_diagonal[-1] == 0)
#         # Solution y1 is the same as the last stage
#         object.__setattr__(
#             self,
#             "ssal",
#             lower_b_sol_equal and diagonal_b_sol_equal and explicit_last_stage,
#         )
#         # Vector field - control product k1 is the same across first/last stages.
#         object.__setattr__(
#             self,
#             "fsal",
#             lower_b_sol_equal and diagonal_b_sol_equal and explicit_first_stage,
#         )


# Path: diffrax/solver/fehlberg2.py
import numpy as np
from ..local_interpolation import ThirdOrderHermitePolynomialInterpolation
from .runge_kutta import AbstractERK, ButcherTableau




_fehlberg2_tableau = ButcherTableau(
    a_lower=(np.array([1 / 2]), np.array([1 / 256, 255 / 256])),
    b_sol=np.array([1 / 512, 255 / 256, 1 / 512]),

    b_error=np.array([-1 / 512, 0, 1 / 512]),