# 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
        # too small.)
        s = final_state.s
        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


# 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

        # too small.)
        s = final_state.s
        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