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


import typing
import jax
import jax.flatten_util as fu
import jax.lax as lax
import jax.numpy as jnp
from typing import Optional, Tuple
from ..custom_types import Array, PyTree, Scalar
            if key.startswith("__") and key.endswith("__"):
                _dict[key] = value
            else:
                _dict[key] = i
                reverse_lookup.append(value)
                i += 1
        _dict["reverse_lookup"] = reverse_lookup
        return super().__new__(cls, name, bases, _dict)

    def __instancecheck__(cls, instance):
        return isinstance(instance, int) or super().__instancecheck__(instance)

    def __getitem__(cls, item):
        return cls.reverse_lookup[item]

    def __len__(cls):
        return len(cls.reverse_lookup)


def _fill_forward(
    last_observed_yi: Array["channels":...], yi: Array["channels":...]  # noqa: F821
) -> Tuple[Array["channels":...], Array["channels":...]]:  # noqa: F821
    yi = jnp.where(jnp.isnan(yi), last_observed_yi, yi)
    return yi, yi


@jax.jit
def fill_forward(
    ys: Array["times", ...],  # noqa: F821
    replace_nans_at_start: Optional[Array[...]] = None,  # noqa: F821

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


# Path: diffrax/misc/misc.py
import typing
import jax
import jax.flatten_util as fu
import jax.lax as lax
import jax.numpy as jnp
from typing import Optional, Tuple
from ..custom_types import Array, PyTree, Scalar

            if key.startswith("__") and key.endswith("__"):
                _dict[key] = value
            else:
                _dict[key] = i
                reverse_lookup.append(value)
                i += 1
        _dict["reverse_lookup"] = reverse_lookup
        return super().__new__(cls, name, bases, _dict)

    def __instancecheck__(cls, instance):
        return isinstance(instance, int) or super().__instancecheck__(instance)

    def __getitem__(cls, item):
        return cls.reverse_lookup[item]

    def __len__(cls):
        return len(cls.reverse_lookup)


def _fill_forward(
    last_observed_yi: Array["channels":...], yi: Array["channels":...]  # noqa: F821
) -> Tuple[Array["channels":...], Array["channels":...]]:  # noqa: F821
    yi = jnp.where(jnp.isnan(yi), last_observed_yi, yi)
    return yi, yi


@jax.jit
def fill_forward(
    ys: Array["times", ...],  # noqa: F821
    replace_nans_at_start: Optional[Array[...]] = None,  # noqa: F821

) -> Array["times", ...]:  # noqa: F821