# Path: diffrax/misc/omega.py
# class ω(metaclass=_Metaω):
#     """Provides friendlier syntax for mapping with `jax.tree_map`.
# 
#     !!! example
# 
#         ```python
#         (ω(a) + ω(b)).ω == jax.tree_map(operator.add, a, b)
#         ```
# 
#     !!! tip
# 
#         To minimise the number of brackets, the following `__rpow__` syntax can be
#         used:
# 
#         ```python
#         (a**ω + b**ω).ω == jax.tree_map(operator.add, a, b)
#         ```
# 
#         This is entirely equivalent to the above.
#     """
# 
#     def __init__(self, value, is_leaf=None):
#         """
#         **Arguments:**
# 
#         - `value`: The PyTree to wrap.
#         - `is_leaf`: An optional value for the `is_leaf` argument to `jax.tree_map`.
# 
#         !!! note
# 
#             The `is_leaf` argument cannot be set when using the `__rpow__` syntax for
#             initialisation.
#         """
#         self.ω = value
#         self.is_leaf = is_leaf
# 
#     def __getitem__(self, item):
#         return ω(
#             jax.tree_map(lambda x: x[item], self.ω, is_leaf=self.is_leaf),
#             is_leaf=self.is_leaf,
#         )
# 
#     def call(self, fn):
#         return ω(jax.tree_map(fn, self.ω, is_leaf=self.is_leaf), is_leaf=self.is_leaf)
# 
#     @property
#     def at(self):
#         return _ωUpdateHelper(self.ω, self.is_leaf)


import math
import operator
import diffrax
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jrandom
import pytest
import scipy.stats
from diffrax.misc import ω
from helpers import all_ode_solvers, random_pytree, shaped_allclose, treedefs
                y0,
                stepsize_controller=diffrax.ConstantStepSize(compile_steps=True),
            )
        )(_t1)
    )()
    assert shaped_allclose(sol.stats["compiled_num_steps"], jnp.array([20, 20]))

    _dt0 = jnp.array([0.1, 0.05])
    sol = jax.jit(
        lambda: jax.vmap(
            lambda dt0: diffeqsolve_nojit(
                terms,
                solver,
                0,
                1,
                dt0,
                y0,
                stepsize_controller=diffrax.ConstantStepSize(compile_steps=True),
            )
        )(_dt0)
    )()
    assert shaped_allclose(sol.stats["compiled_num_steps"], jnp.array([20, 20]))


@pytest.mark.parametrize("solver", [diffrax.ImplicitEuler(), diffrax.Kvaerno5()])
def test_grad_implicit_solve(solver):
    # Check that we work around JAX issue #9374
    # Whilst we're at -- for efficiency -- check the use of PyTree-valued state with
    # implicit solves.


# Path: diffrax/misc/omega.py
# class ω(metaclass=_Metaω):
#     """Provides friendlier syntax for mapping with `jax.tree_map`.
# 
#     !!! example
# 
#         ```python
#         (ω(a) + ω(b)).ω == jax.tree_map(operator.add, a, b)
#         ```
# 
#     !!! tip
# 
#         To minimise the number of brackets, the following `__rpow__` syntax can be
#         used:
# 
#         ```python
#         (a**ω + b**ω).ω == jax.tree_map(operator.add, a, b)
#         ```
# 
#         This is entirely equivalent to the above.
#     """
# 
#     def __init__(self, value, is_leaf=None):
#         """
#         **Arguments:**
# 
#         - `value`: The PyTree to wrap.
#         - `is_leaf`: An optional value for the `is_leaf` argument to `jax.tree_map`.
# 
#         !!! note
# 
#             The `is_leaf` argument cannot be set when using the `__rpow__` syntax for
#             initialisation.
#         """
#         self.ω = value
#         self.is_leaf = is_leaf
# 
#     def __getitem__(self, item):
#         return ω(
#             jax.tree_map(lambda x: x[item], self.ω, is_leaf=self.is_leaf),
#             is_leaf=self.is_leaf,
#         )
# 
#     def call(self, fn):
#         return ω(jax.tree_map(fn, self.ω, is_leaf=self.is_leaf), is_leaf=self.is_leaf)
# 
#     @property
#     def at(self):
#         return _ωUpdateHelper(self.ω, self.is_leaf)


# Path: test/test_integrate.py
import math
import operator
import diffrax
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jrandom
import pytest
import scipy.stats
from diffrax.misc import ω
from helpers import all_ode_solvers, random_pytree, shaped_allclose, treedefs

                y0,
                stepsize_controller=diffrax.ConstantStepSize(compile_steps=True),
            )
        )(_t1)
    )()
    assert shaped_allclose(sol.stats["compiled_num_steps"], jnp.array([20, 20]))

    _dt0 = jnp.array([0.1, 0.05])
    sol = jax.jit(
        lambda: jax.vmap(
            lambda dt0: diffeqsolve_nojit(
                terms,
                solver,
                0,
                1,
                dt0,
                y0,
                stepsize_controller=diffrax.ConstantStepSize(compile_steps=True),
            )
        )(_dt0)
    )()
    assert shaped_allclose(sol.stats["compiled_num_steps"], jnp.array([20, 20]))


@pytest.mark.parametrize("solver", [diffrax.ImplicitEuler(), diffrax.Kvaerno5()])
def test_grad_implicit_solve(solver):
    # Check that we work around JAX issue #9374
    # Whilst we're at -- for efficiency -- check the use of PyTree-valued state with
    # implicit solves.


    term = diffrax.ODETerm(lambda t, y, args: (-args * y**ω).ω)