# 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/custom_types.py
# class PyTree:
#     def __class_getitem__(cls, item):
#         class X:
#             pass
# 
#         X.__module__ = "builtins"
#         X.__qualname__ = _maybe_tuple_to_str(item)
#         return _PyTree[X]


import operator
import equinox as eqx
import jax
import jax.numpy as jnp
from ..brownian import AbstractBrownianPath
from ..custom_types import PyTree
    diffusion: callable

    def __call__(self, t, y, args):
        y, _ = y
        diffusion = self.diffusion(t, y, args)
        return diffusion, 0.0


class _AugBrownianPath(eqx.Module):
    bm: AbstractBrownianPath

    @property
    def t0(self):
        return self.bm.t0

    @property
    def t1(self):
        return self.bm.t1

    def evaluate(self, t0, t1):
        return self.bm.evaluate(t0, t1), 0.0


def sde_kl_divergence(
    *,
    drift1: callable,
    drift2: callable,
    diffusion: callable,
    context: callable,
    y0: PyTree,

# 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/custom_types.py
# class PyTree:
#     def __class_getitem__(cls, item):
#         class X:
#             pass
# 
#         X.__module__ = "builtins"
#         X.__qualname__ = _maybe_tuple_to_str(item)
#         return _PyTree[X]


# Path: diffrax/misc/sde_kl_divergence.py
import operator
import equinox as eqx
import jax
import jax.numpy as jnp
from ..brownian import AbstractBrownianPath
from ..custom_types import PyTree

    diffusion: callable

    def __call__(self, t, y, args):
        y, _ = y
        diffusion = self.diffusion(t, y, args)
        return diffusion, 0.0


class _AugBrownianPath(eqx.Module):
    bm: AbstractBrownianPath

    @property
    def t0(self):
        return self.bm.t0

    @property
    def t1(self):
        return self.bm.t1

    def evaluate(self, t0, t1):
        return self.bm.evaluate(t0, t1), 0.0


def sde_kl_divergence(
    *,
    drift1: callable,
    drift2: callable,
    diffusion: callable,
    context: callable,
    y0: PyTree,

    bm: AbstractBrownianPath,