# Path: diffrax/custom_types.py
# class Array:
#     def __class_getitem__(cls, item):
#         class X:
#             pass
# 
#         X.__module__ = "builtins"
#         X.__qualname__ = _maybe_tuple_to_str(item)
#         return _Array[X]
# 
# class Scalar:
#     pass
# 
# 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/brownian/base.py
# class AbstractBrownianPath(AbstractPath):
#     "Abstract base class for all Brownian paths."
# 
#     @abc.abstractmethod
#     def evaluate(self, t0: Scalar, t1: Scalar, left: bool = True) -> Array:
#         r"""Samples a Brownian increment $w(t_1) - w(t_0)$.
# 
#         Each increment has distribution $\mathcal{N}(0, t_1 - t_0)$.
# 
#         **Arguments:**
# 
#         - `t0`: Start of interval.
#         - `t1`: End of interval.
#         - `left`: Ignored. (This determines whether to treat the path as
#             left-continuous or right-continuous at any jump points, but Brownian
#             motion has no jump points.)
# 
#         **Returns:**
# 
#         A JAX array corresponding to the increment $w(t_1) - w(t_0)$.
# 
#         Some subclasses may allow `t1=None`, in which case just the value $w(t_0)$ is
#         returned.
#         """


from dataclasses import field
from typing import Optional, Tuple
from ..custom_types import Array, Scalar
from ..misc import error_if
from .base import AbstractBrownianPath
import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
import jax.random as jrandom
        u = final_state.u
        w_s = final_state.w_s
        w_t = final_state.w_t
        w_u = final_state.w_u
        rescaled_τ = (τ - s) / (u - s)
        # Fit polynomial as usual.
        # The polynomial ax^2 + bx + c is found by solving
        # [s^2 s 1][a]   [w_s]
        # [t^2 t 1][b] = [w_t]
        # [u^2 u 1][c]   [w_u]
        #
        # `A` is inverse of the above matrix, rescaled to s=0, t=0.5, u=1.
        A = jnp.array([[2, -4, 2], [-3, 4, -1], [1, 0, 0]])
        coeffs = jnp.tensordot(A, jnp.stack([w_s, w_t, w_u]), axes=1)
        return jnp.polyval(coeffs, rescaled_τ)


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

- `t0`: The start of the interval the Brownian motion is defined over.
- `t1`: The start of the interval the Brownian motion is defined over.
- `tol`: The discretisation that `[t0, t1]` is discretised to.
- `shape`: What shape each individual Brownian sample should be.
- `key`: A random key.

!!! info

    If using this as part of an SDE solver, and you know (or have an estimate of) the
    step sizes made in the solver, then you can optimise the computational efficiency

# Path: diffrax/custom_types.py
# class Array:
#     def __class_getitem__(cls, item):
#         class X:
#             pass
# 
#         X.__module__ = "builtins"
#         X.__qualname__ = _maybe_tuple_to_str(item)
#         return _Array[X]
# 
# class Scalar:
#     pass
# 
# 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/brownian/base.py
# class AbstractBrownianPath(AbstractPath):
#     "Abstract base class for all Brownian paths."
# 
#     @abc.abstractmethod
#     def evaluate(self, t0: Scalar, t1: Scalar, left: bool = True) -> Array:
#         r"""Samples a Brownian increment $w(t_1) - w(t_0)$.
# 
#         Each increment has distribution $\mathcal{N}(0, t_1 - t_0)$.
# 
#         **Arguments:**
# 
#         - `t0`: Start of interval.
#         - `t1`: End of interval.
#         - `left`: Ignored. (This determines whether to treat the path as
#             left-continuous or right-continuous at any jump points, but Brownian
#             motion has no jump points.)
# 
#         **Returns:**
# 
#         A JAX array corresponding to the increment $w(t_1) - w(t_0)$.
# 
#         Some subclasses may allow `t1=None`, in which case just the value $w(t_0)$ is
#         returned.
#         """


# Path: diffrax/brownian/tree.py
from dataclasses import field
from typing import Optional, Tuple
from ..custom_types import Array, Scalar
from ..misc import error_if
from .base import AbstractBrownianPath
import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
import jax.random as jrandom

        u = final_state.u
        w_s = final_state.w_s
        w_t = final_state.w_t
        w_u = final_state.w_u
        rescaled_τ = (τ - s) / (u - s)
        # Fit polynomial as usual.
        # The polynomial ax^2 + bx + c is found by solving
        # [s^2 s 1][a]   [w_s]
        # [t^2 t 1][b] = [w_t]
        # [u^2 u 1][c]   [w_u]
        #
        # `A` is inverse of the above matrix, rescaled to s=0, t=0.5, u=1.
        A = jnp.array([[2, -4, 2], [-3, 4, -1], [1, 0, 0]])
        coeffs = jnp.tensordot(A, jnp.stack([w_s, w_t, w_u]), axes=1)
        return jnp.polyval(coeffs, rescaled_τ)


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

- `t0`: The start of the interval the Brownian motion is defined over.
- `t1`: The start of the interval the Brownian motion is defined over.
- `tol`: The discretisation that `[t0, t1]` is discretised to.
- `shape`: What shape each individual Brownian sample should be.
- `key`: A random key.

!!! info

    If using this as part of an SDE solver, and you know (or have an estimate of) the
    step sizes made in the solver, then you can optimise the computational efficiency

    of the Virtual Brownian Tree by setting `tol` to be just slightly smaller than the