packquickly

import functools as ft
from collections.abc import Callable

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.tree_util as jtu
import lineax as lx
import optimistix as optx
from equinox.internal import ω

packquickly

import equinox as eqx  # https://github.com/patrick-kidger/equinox
import jax
import jax.numpy as jnp
import optax
import optimistix as optx


# HIMMELBG test problem from CUTE
@eqx.filter_jit
def toy_problem(y1, y2, constants):