Detected differentiation of a custom_jvp function with respect to a closed-over value?

[adam-hartshorne]

I was trying to follow the example in Chapter 2 of the http://implicit-layers-tutorial.org/implicit_functions/ written by @mattjj that looks at ODEs.

I get the following error

Detected differentiation of a custom_jvp function with respect to a closed-over value,

which links to this page, https://jax.readthedocs.io/en/latest/custom_vjp_update.html

But I have no idea really what this means or how you go about updating the examples to fit for this new approach.

In particular I am looking at the same @odeint_rk4.defjvp example.

Any help much appreciated.

[agobL]

I refresh this discussion because I had the same issue.

[mattjj]

I just pushed that updated colab notebook to the implicit-layers-tutorial.org site (I'm impressed that Zico's setup works instantly even after not touching it for a year!).

[agobL]

Your notebook works, but I share you my code so that you can see the error :

from jax import custom_jvp,jvp
import jax.numpy as np
from jax.lax import scan
from noloadj.optimization.optimProblem import OptimProblem,Spec
import math

def odeint_rk4(f, y0, t, *args):
  def step(state, t):
    y_prev, t_prev = state
    h = t - t_prev
    k1 = h * f(y_prev, t_prev, *args)
    k2 = h * f(y_prev + k1/2., t_prev + h/2., *args)
    k3 = h * f(y_prev + k2/2., t_prev + h/2., *args)
    k4 = h * f(y_prev + k3, t + h, *args)
    y = y_prev + 1./6 * (k1 + 2 * k2 + 2 * k3 + k4)
    return (y, t), y
  _, ys = scan(step, (y0, t[0]), t[1:])
  return ys

odeint_rk4 = custom_jvp(odeint_rk4, nondiff_argnums=(0,))

@odeint_rk4.defjvp
def odeint_rk4_jvp(f, primals, tangents):
  y0, t, *args = primals
  delta_y0, _, *delta_args = tangents
  nargs = len(args)

  def f_aug(aug_state, t, *args_and_delta_args):
    primal_state, tangent_state = aug_state
    args, delta_args = args_and_delta_args[:nargs], args_and_delta_args[nargs:]
    primal_dot, tangent_dot = jvp(f, (primal_state, t, *args), (tangent_state, 0., *delta_args))
    return np.stack([primal_dot, tangent_dot])

  aug_init_state = np.stack([y0, delta_y0])
  aug_states = odeint_rk4(f_aug, aug_init_state, t, *args, *delta_args)
  ys, ys_dot = aug_states[:, 0, :], aug_states[:, 1, :]
  return ys, ys_dot


def lancer(m,R,v0,a, x0, y0, tf):
    g = 9.81
    k=0.5*1.292*0.5*math.pi*R*R # k=0.5*rho(air)*Cx*Section_sphere pour f=-kv²
    # frottements  https://www.physagreg.fr/mecanique-12-chute-frottements.php
    vx0,vy0= v0*np.cos(a),v0*np.sin(a)
    sol = odeint_rk4(lambda s,t : np.array([s[2],s[3],-k*s[2]*(s[2]*s[2]+s[3]*s[3])**0.5/m, -k*s[3]*(s[2]*s[2]+s[3]*s[3])**0.5/m-g]),
                 np.array([x0, y0, vx0, vy0]),np.linspace(0,tf,int(tf/1e-3)))
    s=np.transpose(sol)
    x,y=s[0],s[1]
    yf,xf=y[-1],x[-1]
    hauteur=np.max(y)
    return locals().items()

spec=Spec(variables={'m':1.0,'R':0.2,'v0':10,'a':math.pi/4}, bounds={'m':[0.5,5.],'R':[0.01,1.0],'v0':[10.,100.],'a':[0.,math.pi/2]},
          objectives=['hauteur'], eq_cstr={'yf':0.0,'xf':22.0},freeOutputs=['x','y'])

parameters={'x0':0.,'y0':2.,'tf':3}
optim=OptimProblem(model=lancer,specifications=spec,parameters=parameters)
result=optim.run()
result.printResults()

[mattjj]

Ah okay great, so just to make sure I understand:

1. we think the notebook now posted to implicit-layers-tutorial.org now works, and
2. we want to debug the code you just pasted.

Does that sound right?

By the way, there are much more developed ODE differentiation libraries at this point, like Diffrax. Are you trying to do something not supported by Diffrax? (E.g. maybe you want a minimal self-contained implementation, or maybe Diffrax doesn't support forward mode.)

I'll take a look at the code you shared when I can! In the meantime, let me know your thoughts on using Diffrax.

[agobL]

I didn't know Diffrax, thanks for the discovering. Actually, I am a phd student so I prefer to use my own functions to solve ODE. I am focused on dynamic optimization for systems sizing, and most of the existing librairies focus on optimal control.

[agobL]

Any news ?

[YouJiacheng]

Maybe helpful: #10134 (comment)

[agobL]

I found how to manage the derivative in #10134