piecewise function with apparently continuous gradients - but jnp.where produces nan gradients.

[HansN87]

Hi,

I am struggling to understand the following issue:
I have a function of x with parameters a, b that is defined piecewise in two segments , x in [-s,s] and x in (s, infinity), or x in [-s, s) and x in [s, infinity). I can evaluate the gradients in the parameters a, b in both branches separately if x falls in the respective segment. They work just fine. They appear to be continuous across the transition at x=s, verified with two approaches: using jax.grad and using my own finite difference calculation. Yet, if I combined the two functions using jnp.where, I observe nan gradients for x values in the left branch. The same values of x lead to well-behaved gradients in the same left branch if I do not use jnp.where and only evaluate the left segment itself directly.

For my project, however, I need to be able to evaluate both branches and combined in a way that avoids nan grads.
I'd appreciate any input. Asking here first, since I am not sure whether it's an actual bug in jax.

See the notebook below for details.
nan_gradient_with_biweight_kernel_convoluted_gamma_pdf (2).pdf

Background:
The function I work with is a convolution of a gamma pdf (https://en.wikipedia.org/wiki/Gamma_distribution) which is supported on [0, infinity) with a biweight kernel function ( https://en.wikipedia.org/wiki/Kernel_(statistics) ) that is supported on (-s, s). The function itself was obtained through symbolic integration with Mathematica. Unfortunately I don't have access to analytic gradients. But based on my checks with finite differences, I can't find anything wrong with it - and the issue only appears if jnp.where is involved.

Thank you.

[jakevdp]

It sounds like this FAQ entry will be relevant: Gradients contain NaN when using where.

[HansN87]

Thank you. This helped me fixing it. I didn't realize that branch1 was unstable in the domain of branch0.

def combine(x, a, b, sigma):
    s = 3.0 * sigma
    x0 = jnp.where(x < s, x, s)
    b0 = branch0(x0, a, b, sigma)
    
    x1 = jnp.where(x < s, s, x)
    b1 = branch1(x1, a, b, sigma)
    return jnp.where(x < s, b0, b1)

notebook below in case anyone finds this in the future.
nan_gradient_with_biweight_kernel_convoluted_gamma_pdf_v2.pdf