Making diffrax Haiku compliant

Was trying to integrate diffrax diffeqsolve into a Haiku module but was running into similar issues as mentioned in #115 and figured out how to do it and thought I'd share in case others had the same issues or someone wanted to add my solution to the documentation. Although I'm fairly new to Haiku and Diffrax so my solution may not be the only/cleanest.

The main conflicts between Haiku and Diffrax are:

• Haiku transforms don't play nice with other transforms (including equinox filter_jit)
• Native haiku module calls aren't compliant with ODETerm (ie aren't of the form module(ts, y0, args))

The first point can be addressed by simply unwrapping diffeqsolve. The second requires (in my solution) two classes to wrap a given module; one to integrate the unwrapped diffeq_solve into the __call__ of an outer class and other to wrap a hk.Module call to work with ODETerm in the outer odenet

Here's a small example using hk.nets.MLP

import jax
import dataclasses


import haiku as hk
import jax.numpy as jnp

from diffrax import diffeqsolve, ODETerm, Dopri5

#NOTE: need to unwrapp diffeqsolve to include it in a haiku module
orig_diffeqsolve = diffeqsolve.__wrapped__

#wrapper for full networks already defined by haiku
#separate classes should be added for specific networks not available (such as N-BEATS)
@dataclasses.dataclass
class odenet(hk.Module):
    def __init__(self, fn: hk.Module, fn_args: dict):
        super().__init__(fn_args["name"]+"_odeint")
        self.fn = ODETerm(fn(**fn_args))
        #
    def __call__(self, ts, y0) -> jax.Array:
        #diffrax attempt:
        solution = orig_diffeqsolve(
            self.fn,
            Dopri5(),
            t0 = ts[0],
            t1 = ts[1],
            dt0 = 1e-2,
            y0 = y0)
        #
        return solution.ys[0]



#wrapper class needed to define __call__ form compliant with diffeqsolve
class MLP(hk.Module):
    def __init__(self, output_sizes, name):
        super().__init__(name)
        self.fn = hk.nets.MLP(output_sizes)
    def __call__(self, ts, y0, args) ->jax.Array:
        return self.fn(y0)


def forward_pass(ts, y0):
    module = odenet(fn=MLP, fn_args={"output_sizes":[3, 10, 3], "name":"MLP"})
    return module(ts, y0)


dummy_x = jnp.array([[1., 2., 3.]])


fwrd = hk.transform(forward_pass)

key = jax.random.PRNGKey(120)
params = fwrd.init(rng=key,ts=[0,1], y0=dummy_x)

pred = fwrd.apply(params=params, ts=[0, 1.2], y0=dummy_x, rng=key)

There's probably a better way to wrap the lower level hk.Module but figured this was good enough to share

Start with the Haiku/MLP example in this issue and compare it with the related discussion in #115. Document a clear Diffrax–Haiku integration pattern, including the wrapper constraints and a working example, and make the guidance available in the project's documentation.