Model is slower using diffrax, and even more when calling a function inside the vector_field

Hello,

I am using diffrax to time-forward an ODE in a context of variational inversion.
Using equinox fitlers and abstract solver from diffrax is very powerful, and it could be very beneficial for my project down the line.

I coded a simple model in both equinox/diffrax and JAX pure but i see that the former is about 3 times slower.
For now this is negligible but later I will code a heavier model and it will become problematic.

The key aspect of this model is that I run many many timesteps (28 days at dt=60s, 40320 timesteps).
However, it seems that the diffrax model is always about x3 compared to JAX pure no matter the number of iterations.

For the gradient of the cost function, I use forward AD because of my very limited number of control parameters, but using diffrax should allow me to switch easily to reverse AD with automatic checkpointing if needed.

Description of my models

Note : the equation solved is the same for all

• jmodel : JAX only code, my model is a Python class with a 'do_forward' method. Forcing is hard coded in the time loop.
• eqx_model : my model is a eqx.Module class with a 'call' method. Forcing is hard coded in the time loop.
• eqx_model_interpn : my model is a eqx.Module class with a 'call' method. Forcing calls a interpolation function.
• eqx_dfx_model : my model is a eqx.Module class with a 'call' method, and I use diffrax to time step the model. Interpolation is hard coded in the vector_field definition.

Benchmark results (seconds)

forward pass:

jmodel : 0.18528
eqx_model : 0.22648
eqxmodel_interpn : 3.45376
eqx_dfx_model  : 0.53019

cost :

jmodel : 0.18697
eqx_model : 0.23840
eqxmodel_interpn : 3.25355
eqx_dfx_model  : 0.59816

grad :

jmodel : 0.26339
eqx_model : 0.29954
eqxmodel_interpn : 3.10121
eqx_dfx_model  : 1.02204

My questions

• jmodel vs eqx_model: Equinox adds a bit of overhead for each timestep (about 1 microsec/timestep), I think this is normal but could you confirm ? (related to this)
• eqx_model vs eqxmodel_interpn: Evaluating a function inside the timeloop seems to increase a lot the total time for the forward model (about x10 in my example). Is this expected ? I guess I have a workaround with my interpolation 'by hand' but it would be awesome to change the interpolation method with coding this part again, using diffrax.LinearInterpolation.
• eqx_dfx_model vs eqx_model: Is this overhead induced by diffrax considered normal ? If no, what option in diffeqsolve could provide similar performance to pure JAX ? (I have had a look at #517 and #592 but had no success with modifying the diffeqsolve arguments)

Thank you for you help !

Hugo

config:
I run on a T400 mobile GPU
diffrax 0.7
jax 0.5.2
equinox 0.11.12

Code

import numpy as np
import os
import matplotlib.pyplot as plt
import time as clock
import jax
import jax.numpy as jnp
from jax import jit, lax
import jax.tree_util as jtu
from functools import partial
import equinox as eqx
import diffrax
from diffrax import Euler, diffeqsolve, ODETerm
os.environ["EQX_ON_ERROR"] = "nan"
jax.config.update("jax_enable_x64", True)

class classic_slab1D:
    def __init__(self, TAx, TAy, fc, t0, nt, dt, dt_forcing):
        self.TA = jnp.asarray(TAx) + 1j*jnp.asarray(TAy)
        self.fc = jnp.asarray(fc)  
        self.t0 = t0
        self.t1 = t0+nt*dt
        self.dt = dt
        self.dt_forcing = dt_forcing
        self.nt = nt
        self.do_forward_jit = jit(self.do_forward)
        
    def __one_step(self, X0, it):
        K, U = X0      
        # interpolation
        nsubsteps = self.dt_forcing // self.dt
        itf = jnp.array(it//nsubsteps, int)
        aa = jnp.mod(it,nsubsteps)/nsubsteps
        itsup = lax.select(itf+1>=self.nt, -1, itf+1) 
        TA = (1-aa)*self.TA[itf] + aa*self.TA[itsup]
        # 1 time forward
        U = U.at[it+1].set( U[it] + self.dt*(-1j*self.fc*U[it] 
                                           + K[0]*TA 
                                           - K[1]*U[it] ) )  
        X0 = K, U        
        return X0, X0

    def do_forward(self, pk):
        # initialisation
        U = jnp.zeros( self.nt, dtype='complex')
        K = jnp.exp(pk)
        # time loop
        X0 = K, U
        final, _ = lax.scan(self.__one_step, X0, jnp.arange(0,self.nt-1))
        _, U = final
        return jnp.real(U), jnp.imag(U)
   
class classic_slab1D_eqx(eqx.Module):
    
    pk : jnp.array # control vector
    TA : jnp.array          # parameters
    fc : jnp.array          # |
    dt_forcing : np.int32   # |
    t0 : np.int32           # run parameters
    t1 : np.int32           # |
    nt : np.int32           # | 
    dt : np.int32           # |
    
    is_difx : bool          # use diffrax or not
    interp_fn : bool        # use interoplation function or not
    
    def __init__(self, pk, TAx, TAy, fc, t0, nt, dt, dt_forcing, is_difx, interp_fn=False):
        self.fc = jnp.asarray(fc)  
        self.TA = jnp.asarray(TAx) + 1j*jnp.asarray(TAy)
        self.pk = pk
        self.t0 = t0
        self.t1 = t0+nt*dt
        self.dt = dt
        self.dt_forcing = dt_forcing
        self.nt = nt
        self.is_difx = is_difx
        self.interp_fn = interp_fn
        
    @eqx.filter_jit
    def __call__(self):
        K = jnp.exp(self.pk)
        nsubsteps = self.dt_forcing // self.dt
        
        if self.interp_fn:
                TAx_t, TAy_t = self.TA_at_t(jnp.real(self.TA), jnp.imag(self.TA), self.t0, self.t1, self.dt_forcing)
        
        if self.is_difx: # diffrax time forward

            def vector_field(t, C, args):
                U,V = C
                fc, K, TAx, TAy = args
                if self.interp_fn:
                    TAx,TAy = TAx_t.evaluate(t), TAy_t.evaluate(t)
                else:
                    # on the fly interpolation
                    it = jnp.array(t//dt, int)
                    itf = jnp.array(it//nsubsteps, int)
                    aa = jnp.mod(it,nsubsteps)/nsubsteps
                    itsup = lax.select(itf+1>=len(TAx), -1, itf+1) 
                    TAx = (1-aa)*TAx[itf] + aa*TAx[itsup]
                    TAy = (1-aa)*TAy[itf] + aa*TAy[itsup]
                # physic
                d_U = fc*V + K[0]*TAx - K[1]*U
                d_V = -fc*U + K[0]*TAy - K[1]*V
                d_y = d_U,d_V
                return d_y

            sol = diffeqsolve(terms=ODETerm(vector_field), 
                           solver=Euler(), 
                           t0=self.t0, 
                           t1=self.t1, 
                           y0=(0.0, 0.0), 
                           args=(self.fc, K, jnp.real(self.TA), jnp.imag(self.TA)), 
                           dt0=dt, 
                           saveat=diffrax.SaveAt(steps=True),
                           adjoint=diffrax.ForwardMode(),
                           max_steps=nt).ys
            U = sol[0]+1j*sol[1]
            
        else: # my time forward
            
            U = jnp.zeros( self.nt, dtype='complex')
            def __one_step(X0, it):
                K, U = X0  
                # interpolation
                if self.interp_fn:
                    t = self.t0 + it*self.dt
                    TA = TAx_t.evaluate(t) +1j*TAy_t.evaluate(t)
                else:
                    itf = jnp.array(it//nsubsteps, int)
                    aa = jnp.mod(it,nsubsteps)/nsubsteps
                    itsup = lax.select(itf+1>=self.nt, -1, itf+1) 
                    TA = (1-aa)*self.TA[itf] + aa*self.TA[itsup]
                # one time forward
                U = U.at[it+1].set( U[it] + self.dt*(-1j*self.fc*U[it] 
                                                + K[0]*TA 
                                                - K[1]*U[it] ) )  
                X0 = K, U        
                return X0, X0
            # time loop
            X0 = K, U
            final, _ = lax.scan(lambda carry, y: __one_step(carry, y), X0, jnp.arange(0,self.nt-1)) # 
            _, U = final
        return jnp.real(U),jnp.imag(U)
    def TA_at_t(self, TAx, TAy, t0, t1, dt_forcing):
        """
        return a function that interpolate at t the forcing
        """
        time_forcing = jnp.arange(t0, t1, dt_forcing, dtype=float)
        # print(len(time_forcing),len(TAx))
        TAx_t = diffrax.LinearInterpolation(time_forcing, TAx)
        TAy_t = diffrax.LinearInterpolation(time_forcing, TAy)
        return TAx_t, TAy_t

def benchmark(func, N=20):
    L = np.zeros(N)
    _ = func() # run once for compilation
    for k in range(N):
        time1=clock.time()
        result = func()
        L[k] = clock.time()-time1
    return L.mean(), L.std()

def cost(sol, obs): return jnp.nanmean( (sol[0]-obs[0])**2 + (sol[0]-obs[0])**2)
    
def cost_j(pk, jmodel, obs):
    sol = jmodel.do_forward_jit(pk)
    return cost(sol, obs)

def cost_eqx(dynamic_model, static_model, obs):
    mymodel = eqx.combine(dynamic_model, static_model)
    sol = mymodel()
    return cost(sol, obs)

def dcost_j(pk, jmodel, obs): return jax.jacfwd(cost_j)(pk, jmodel, obs)
    
def dcost_eqx(dynamic_model, static_model, obs): return eqx.filter_jacfwd(cost_eqx)(dynamic_model, static_model, obs)
    
def my_partition(mymodel):
    filter_spec = jtu.tree_map(lambda arr: False, mymodel) # keep nothing
    filter_spec = eqx.tree_at( lambda tree: tree.pk, filter_spec, replace=True) # keep only pk
    return eqx.partition(mymodel, filter_spec)
    
# MY PARAMETERS
t0 = 0
nt = 28*86400//60
dt = 60
dt_forcing = 3600.
time_forcing = jnp.arange(t0,nt*dt,dt_forcing)
TAx = 0.2*np.ones(len(time_forcing)) # step forcing
TAy = 0.0*np.ones(len(time_forcing))
fc = 1e-4 

pktarget = jnp.asarray([-8.,-13.])
pkini = jnp.array([-9.,-11.])
 

jmodel = classic_slab1D(TAx, TAy, fc, t0, nt, dt, dt_forcing)
eqxmodel = classic_slab1D_eqx(pktarget, TAx, TAy, fc, t0, nt, dt, dt_forcing, is_difx=False)
eqxmodel_dfx = classic_slab1D_eqx(pktarget, TAx, TAy, fc, t0, nt, dt, dt_forcing, is_difx=True)
eqxmodel_fninterp = classic_slab1D_eqx(pktarget, TAx, TAy, fc, t0, nt, dt, dt_forcing, is_difx=False, interp_fn=True)
N = 10

# make some observations
dt_obs = 86400 # 1 per day
obs = np.nan*jnp.zeros((2,nt),dtype=np.float64)
truth = jmodel.do_forward_jit(pktarget)
for k in range(nt):
    if k%(dt_obs//dt)==0:
        obs = obs.at[0,k].set(truth[0][k]) # U
        obs = obs.at[1,k].set(truth[1][k]) # V

# new models with new control parameter pk
eqxmodel_2 = classic_slab1D_eqx(pkini, TAx, TAy, fc, t0, nt, dt, dt_forcing, is_difx=False)
eqxmodel_fninterp_2 = classic_slab1D_eqx(pktarget, TAx, TAy, fc, t0, nt, dt, dt_forcing, is_difx=False, interp_fn=True)
eqxmodel_dfx_2 = classic_slab1D_eqx(pkini, TAx, TAy, fc, t0, nt, dt, dt_forcing, is_difx=True)
dyn_eqx, stat_eqx = my_partition(eqxmodel_2)
dyn_eqx_fni, stat_eqx_fni = my_partition(eqxmodel_fninterp_2)
dyn_eqx_dfx, stat_eqx_dfx = my_partition(eqxmodel_dfx_2)

print('Forward model:')
print('     jmodel:           mean, std (s)', benchmark(partial(jmodel.do_forward_jit,pk=pktarget),N=N))
print('     eqxmodel:         mean, std (s)', benchmark(eqxmodel,N=N))
print('     eqxmodel_interpn: mean, std (s)', benchmark(eqxmodel_fninterp,N=N))
print('     eqx_dfx_model:    mean, std (s)', benchmark(eqxmodel_dfx,N=N))

print('Cost:')
print('     jmodel:           mean, std (s)', benchmark(partial(cost_j,pk=pkini,jmodel=jmodel,obs=obs),N=N))
print('     eqxmodel:         mean, std (s)', benchmark(partial(cost_eqx,dyn_eqx,stat_eqx,obs),N=N))
print('     eqxmodel_interpn: mean, std (s)', benchmark(partial(cost_eqx,dyn_eqx_fni,stat_eqx_fni,obs),N=N))
print('     eqx_dfx_model:    mean, std (s)', benchmark(partial(cost_eqx,dyn_eqx_dfx,stat_eqx_dfx,obs),N=N))

print('gradient (forward):')
print('     jmodel:           mean, std (s)', benchmark(partial(dcost_j,pk=pkini,jmodel=jmodel,obs=obs),N=N))
print('     eqxmodel:         mean, std (s)', benchmark(partial(dcost_eqx,dyn_eqx,stat_eqx,obs),N=N))
print('     eqxmodel_interpn: mean, std (s)', benchmark(partial(dcost_eqx,dyn_eqx_fni,stat_eqx_fni,obs),N=N))
print('     eqx_dfx_model:    mean, std (s)', benchmark(partial(dcost_eqx,dyn_eqx_dfx,stat_eqx_dfx,obs),N=N))

Reproduce the supplied benchmark using the shown diffeqsolve, ODETerm, Euler, and LinearInterpolation entry points under the listed Python, JAX, Equinox, and diffrax versions. Compare the pure JAX, Equinox, interpolation, and diffrax variants on the stated GPU; done means the overhead is explained and any relevant configuration or limitation is documented.