Advice on what to do in situations when getting inf from jnp.exp

[mjhoover1]

Hello, I have a function that calculates the sigmoid taking in values x and an x_offset. Note that this error only occurs when using jax.jacfwd, e.g., calling the function by itself has no error. Here is the function:

@jit
def sigmoid2(x, x_offset):
    sigmoid_slope = 32.
    return 1. / (1. + jnp.exp(-sigmoid_slope * (x - x_offset)))

The input x is ConcreteArray([-258.614286 0. 40. 80. ] and x_offset is -.1 the following calculation computes as follows:

(x - x_offset) = ConcreteArray([-2.58514286e+02 1.00000000e-01 4.01000000e+01 8.01000000e+01] implies

-sigmoid_slope * (x-x_offset) = ConcreteArray([ 8.27245715e+03 -3.20000000e+00 -1.28320000e+03 -2.56320000e+03] implies

1 + jnp.exp(-sigmoid_slope * (x - x_offset)) = ConcreteArray([ inf 1.0407622 1. 1. ]

So we have an inf showing up here usually
1. / (1. + jnp.exp(-sigmoid_slope * (x - x_offset)))
raises the error FloatingPointError invalid value (nan) encountered in mul since I'm using jacfwd this is due to this calculation causing a Nan in the grad computation. I searched for a solution, but the solutions given seem to be a specific solution for their specific problem like reformulating the equation. However, I don't think it's possible to reformulate the above equation. I looked at the Sharp bits page, but it seems to only give the possible reason for not using config.update("jax_enable_x64", True). But, I am using this as well as config.update("jax_debug_nans", True) and config.update('jax_disable_jit', True) which help me debug the issue. Any advice on how to go about solving this problem without changing the functionality of the function, thanks!

[jakevdp]

The main way to deal with this kind of thing is to define a custom JVP. For example, the jax.scipy.special.expit function is essentially equivalent to a sigmoid, and deals with NaN issue this way: https://github.com/google/jax/blob/4e219220558a2279de2abc9dc4140dedeb61703f/jax/_src/scipy/special.py#L98-L104

Even easier would be to rewrite your function in terms of expit, so that you can take advantage of this custom JVP without having to re-implement it yourself.

from jax.scipy.special import expit
def sigmoid2(x, x_offset):
    sigmoid_slope = 32.
    return expit(sigmoid_slope * (x - x_offset))

[mjhoover1]

Works like a charm, thanks @jakevdp!