Here is a proposal for defining a model whose parameters are set at initialisation and which can then be called with the expected signatures for other methods.
from diffrax import diffeqsolve, ODETerm, Dopri5
import jax.numpy as jnp
def f(t, y, args):
return -y
term = ODETerm(f)
solver = Dopri5()
y0 = jnp.array([2., 3.])
solution = diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=y0)
# Dataclass version
from dataclasses import dataclass
# frozen: makes the objets immutable after creation
# so parameters can not be modified at runtime
# it also makes the class hashable, as required by Equinox:
# ValueError: Non-hashable static arguments are not supported.
# kw_only: require the parameter names if they want
# to be set when the object is created
@dataclass(frozen=True, kw_only=True)
class Model:
a:float = 1.0
def __call__(self, t, y, args):
return -self.a*y
m = Model(a=2.0)
sol = diffeqsolve(ODETerm(m), solver, t0=0, t1=1, dt0=0.1, y0=y0)
# can also pass the complete or partial parameters from a dict
# params = {"a": 2.0}
# m = Model(**params)
# no performance difference
# jax.jit seems to have no effect
# term = ODETerm(f)
# term = ODETerm(Model())
# term = ODETerm(jax.jit(Model()))
# %timeit sol = diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=y0)