Compatibility with Flax

Will diffrax.diffeqsolve work inside a Flax linen Module? How would you set up the initialization to use Flax inside of ODETerm instead of Equinox?

[patrick-kidger]

Hurrah! I’m glad you figured this out.

[ameya98]

Ah, I figured out a (slightly hacky) way to do this:

class NeuralODE(flax.struct.PyTreeNode):
    """A simple neural ODE."""

    encoder: nn.Module
    derivative_net: nn.Module
    decoder: nn.Module

    def init(self, rng, coords):
        rng, encoder_rng, derivative_net_rng, decoder_rng = jax.random.split(rng, 4)
        coords, encoder_params = self.encoder.init_with_output(encoder_rng, coords)
        coords, derivative_net_params = self.derivative_net.init_with_output(derivative_net_rng, coords)
        coords, decoder_params = self.decoder.init_with_output(decoder_rng, coords)

        return {
            "encoder": encoder_params,
            "derivative_net": derivative_net_params,
            "decoder": decoder_params
        }

    def apply(self, params, coords):
        coords = self.encoder.apply(params["encoder"], coords)

        def f(t, y, args):
            return self.derivative_net.apply(params["derivative_net"], y)

        term = diffrax.ODETerm(f)
        solver = diffrax.Euler()
        solution = diffrax.diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=coords)
        coords = solution.ys
        coords = self.decoder.apply(params["decoder"], coords)
        return coords

rng = jax.random.PRNGKey(0)
coords = jnp.ones((1, 4))

model = NeuralODE(
    encoder=nn.Dense(10),
    derivative_net=nn.Dense(10),
    decoder=nn.Dense(4))
params = jax.jit(model.init)(rng, coords)

Then you can simply use this like any other nn.Module:

@jax.jit
def compute_loss(params, coords, true_coords):
    preds = model.apply(params, coords)
    return jnp.abs(preds - true_coords).sum()

grads = jax.grad(compute_loss)(params, coords, jnp.zeros_like(coords))

This just uses flax.struct.PyTreeNode instead of eqx.Module. I didn’t want to mix both of them in my codebase. Thanks a lot for the help!