# 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]
# 
# Path: diffrax/misc/unvmap.py
# def unvmap_any(x):
#     return _unvmap_any_p.bind(x)


import math
import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
from ..custom_types import Array
from .unvmap import unvmap_any


class _InplaceUpdate(eqx.Module):
    pred: Array[bool]

    def __call__(self, val: Array):
        return _InplaceUpdateInner(self.pred, val)

    def merge(self, other: "_InplaceUpdate") -> "_InplaceUpdate":
        return _InplaceUpdate(self.pred & other.pred)


class _InplaceUpdateInner(eqx.Module):
    pred: Array[bool]
    val: Array

    @property
    def at(self):
        return _InplaceUpdateInnerInner(self.pred, self.val)


class _InplaceUpdateInnerInner(eqx.Module):
    pred: Array[bool]
    val: Array

    def __getitem__(self, index: Array):
        return _InplaceUpdateInnerInnerInner(self.pred, self.val, index)


class _InplaceUpdateInnerInnerInner(eqx.Module):

# 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]
# 
# Path: diffrax/misc/unvmap.py
# def unvmap_any(x):
#     return _unvmap_any_p.bind(x)


# Path: diffrax/misc/bounded_while_loop.py
import math
import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
from ..custom_types import Array
from .unvmap import unvmap_any



class _InplaceUpdate(eqx.Module):
    pred: Array[bool]

    def __call__(self, val: Array):
        return _InplaceUpdateInner(self.pred, val)

    def merge(self, other: "_InplaceUpdate") -> "_InplaceUpdate":
        return _InplaceUpdate(self.pred & other.pred)


class _InplaceUpdateInner(eqx.Module):
    pred: Array[bool]
    val: Array

    @property
    def at(self):
        return _InplaceUpdateInnerInner(self.pred, self.val)


class _InplaceUpdateInnerInner(eqx.Module):
    pred: Array[bool]
    val: Array

    def __getitem__(self, index: Array):
        return _InplaceUpdateInnerInnerInner(self.pred, self.val, index)


class _InplaceUpdateInnerInnerInner(eqx.Module):

    pred: Array[bool]