A grad on an integration method leads to NaN... Why?

[jecampagne]

Hello,
Here is a simple Simpson integration routine:

@partial(jit, static_argnums=(0,3))  # f et N ne seront des vecteurs
def jax_simps(f, a,b, N):
    dx = (b - a) / N
    x = jnp.linspace(a, b, N + 1)
    y = f(x)
    w = jnp.ones_like(y)
    w = w.at[2:-1:2].set(2.)   
    w = w.at[1::2].set(4.)
    S = dx / 3. * jnp.einsum('i...,i...',w,y)
    return S

It works quite well and for instance

f5 = lambda x: jnp.sqrt(x) * jnp.exp(-x)
intf5 = lambda x: jax_simps(f5,0.,x,N=2**10)
vintf5 = vmap(intf5)

x = jnp.arange(0,10,0.1)
plt.plot(x,f5(x),label=r"$f(x)=\sqrt{x}e^{-x}$")
plt.plot(x,vintf5(x), label=r"$\int_0^x f(z)dz$")
plt.hlines(np.sqrt(np.pi)/2,x.min(),x.max(),colors='k',ls='--',label=r"$\sqrt{\pi}/2$")
plt.legend();

gives

As an exercise I would like to show that

\frac{d}{dx} \int_0^x f(z) dz = f(x)

but

grad(intf5)(1.)     # leads to DeviceArray(nan, dtype=float64, weak_type=True)

Any help would be welcome. Thanks

[jakevdp]

I did some debugging: using jax.debug_nans, it looks like the problem comes from the backward pass over linspace, and I can reproduce the issue more compactly like this:

def f(x):
  return jnp.sqrt(jnp.linspace(0, x, 5))

jax.jacrev(f)(1.0)
# DeviceArray([nan, nan, nan, nan, nan], dtype=float32, weak_type=True)

Narrowing it down more, the issue seems to come from the first value in the linspace, which is computed with something like this:

def f(x):
  return jnp.sqrt(x * 0.0)

jax.grad(f)(1.0) # NaN

Symbolically, the derivative of sqrt(a * x) is equal to a / sqrt(a * x), which is 0/0, which is NaN.

So it looks like the issue is with your particular f is not differentiable at 0.0. So you could fix this by making the lower bound some small nonzero value, like

intf5 = lambda x: jax_simps(f5,1E-7,x,N=2**10)

[jecampagne]

My answer below is exactly like that ;) but you have been very rapid to answer. thanks.

[jecampagne]

Ho I just realize something that explain the NaN. Here is the explanation. Remember that Simpson integration can be written as

I(a,b,n) = (b-a)/(6n) \sum_{i=0}^{2n} w_i f(x_i)

with

x_i = a + (b-a)/(2n)
w_i = w_0=1=w_{2n}, w_{2k-1}=4 (k=1,...;n), w_{2k}=2 (k:1,...,n-1) 

So, when differentiating wrt b this leads to

dI/db = 1/6n sum_{i=0}^{2n} w_i f(x_i) + (b-a)/(6n) sum_{i=0}^{2n} w_i f'(x_i)  i/(2n)

But in case if f(x)=Sqrt(x)Exp(-x) then f'(0) is not defined so x_0 = a yields the singularity when a is set to 0...

So the trick is to avoid the singularity:

intf5 = lambda x: jax_simps(f5,1e-7,x,N=2**10)
grad(intf5)(1.0)

now gives

DeviceArray(0.36787583, dtype=float64, weak_type=True)

as expected!