neural-ode compatibility with dm-haiku models ?

I am currently following the Neural-ODE tutorial on the Diffrax documentation (https://docs.kidger.site/diffrax/examples/neural_ode/). I was wondering if instead of instantiating Equinox modules the Func and NeuralODE there is a work around to incorporate Haiku modules (hk.Module) ?

My end goal would be to use diffrax.diffeqsolve on a ode term that uses hk.Module as such :

solution = diffrax.diffeqsolve(
            diffrax.ODETerm(self.func),
            diffrax.Tsit5(),
            t0=ts[0],
            t1=ts[-1],
            dt0=ts[1] - ts[0],
            y0=y0,
            stepsize_controller=diffrax.PIDController(rtol=1e-3, atol=1e-6),
            saveat=diffrax.SaveAt(ts=ts),  # saving all of the intermediate states (disc then opt)
        )

where self.func = Func() and Func is a Haiku module as below :

class Func(hk.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.layers = []
        for out_shape in [2,10]:
            self.layers.append(hk.Linear(out_shape, name=name))
            self.layers.append(jax.nn.relu)
        self.layers.append(hk.Linear(1, name=name))

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

[patrick-kidger]

That’s correct. I don’t think you’re doing anything like that at the moment. So you’ll need to pass those params through to the vector field each time you call it.

i.e. something like

class Network(hk.Module):
    ...

    def __call__(self, t, y, args):
        ...

network = hk.without_apply_rng(hk.transform(...))
params = network.init(...)
vector_field = functools.partial(network.apply, params)
diffeqsolve(ODETerm(vector_field), ...)

This issue isn’t anything to do with Diffrax I’m afraid; what you’re describing are difficulties with Haiku. If you want to use Haiku then I suggest reading up on how to use Haiku, and perhaps trying it in a non-neural-ODE context first to gain familiarity with that library.

[thibmonsel]

Im not sure to understand, in order to use transformed_object.apply we need the params of our model that is given by invoking the transformed_object.init