Slow runtime when using jacrev

[noeyusiohc]

Hi

I have a problem with jacrev().
I put the solution of ODE in jacrev and it takes more than a hour.
Is it normal?
When I run this code first, it took only few second,,

def model(g, v, i):
    k = 4
    num = k * g * (r + v - 1 - (r - 1) * g**i)
    denom = i * (r - 1) * (g**(i - 1)) + k * (1 - r - v) * (r + v - 1 - (r - 1) * g**i)
    dfdv = num / denom
    return dfdv

g0 = jnp.array([1.0])
v_vals = np.linspace(0, 1, 100)  

def f(i):
  sol = jax.experimental.ode.odeint(model, g0, v_vals, i)[:,0]
  return sol

i = 0.05

df = jax.jacrev(f)
result = df(i)

print(result)

[jakevdp]

Your code doesn't define r. When I define it as r = 0.0, I can reproduce the slow runtime you're seeing. It seems like your problem is one for which the ODE solver converges very slowly. It's not clear to me why this would be the case (perhaps the problem is ill-conditioned?), but you might be able to work around it by setting a looser tolerance in odeint.

Note that jax.experimental.ode is experimental code that is more-or-less unmaintained at this point; you should only rely on it to the extent that it is useful to you.

[noeyusiohc]

Thanks! :) I'll try changing odeint condition

[patrick-kidger]

Consider using Diffrax instead, which is much more featureful than jax.experimental.ode.

Also remember to place a JIT on the outside of your jacrev. (I'm not sure how much this will help here, as this is another one inside, but it is good practice.)

I'd have to stare at the vector field a little longer but I am also a little antsy that it isn't obviously Lipschitz. That can produce poor numerical behaviour.