Piecewise linear interpolation function gives different return values with and without @jit

[dongxulee]

I want to do a piecewise linear interpolation using the following function definition. With and without @jit function decorator would end up with different results, when I apply the same input value. I really don't get why adding @jit will change the function.

v[:,jnp.array(xpp[:,1], dtype = int),
jnp.array(xpp[:,2], dtype = int),
jnp.array(xpp[:,3], dtype = int),
jnp.array(xpp[:,4], dtype = int),
jnp.array(xpp[:,5], dtype = int)].T

This part helps me to select a list of one-dimensional vectors from v matrix. Which defines the target functions, and I want to interpolate the value of x.

f(ws) = v_1 interpolate value f(x_1)
f(ws) = v_2 interpolate value f(x_2)
f(ws) = v_3 interpolate value f(x_3)
...
f(ws) = v_n. interpolate value f(x_n)

The following is the funciton.

def fit(v, xpp):
    value = vmap(partial(jnp.interp,xp = ws))(x = xpp[:,0], 
                                                  fp = v[:,jnp.array(xpp[:,1], dtype = int),
                                                            jnp.array(xpp[:,2], dtype = int),
                                                            jnp.array(xpp[:,3], dtype = int),
                                                            jnp.array(xpp[:,4], dtype = int),
                                                            jnp.array(xpp[:,5], dtype = int)].T)

Any idea helps, thanks a lot!

[jakevdp]

Can you edit your question to include a minimal reproducible example of the issue? I'd like to help investigate the issue, but I don't know what v, xpp, or ws are meant to contain.

[dongxulee]

Thanks for helping me out, the following example is the best I could do, let me know if it make sense:

Setting the variable:

from jax import partial, vmap, jit
import jax.numpy as jnp
xxx = jnp.array([  0.      ,   0.831025,   3.3241  ,   7.479224,  13.2964  ,
              20.775625,  29.916897,  40.720222,  53.1856  ,  67.31303 ,
              83.1025  , 100.55403 , 119.66759 , 140.44324 , 162.88089 ,
             186.98064 , 212.7424  , 240.16623 , 269.2521  , 300.00003 ])

fff = jnp.array([[ 3.61626554e+00,  3.61627531e+00,            -jnp.inf,
                   -jnp.inf,            -jnp.inf,            -jnp.inf,
        -7.35202107e+37, -7.34325502e+37, -3.50007407e+37,
        -1.64472493e+37, -8.36250276e+36, -2.57747340e+36,
        -5.69342370e+35, -8.81759238e+34, -1.08298795e+34,
        -1.32660759e+32, -2.15171236e+30,  3.61739659e+00,
         3.61751652e+00,  3.61758924e+00]])

x = jnp.array([[2.3543993e+02, 2.0000000e+01, 0.0000000e+00, 0.0000000e+00,
              1.0000000e+00, 1.0000000e+00, 1.1188778e-02]])

Then define a function with @jit, and calculate the fitting value and assign the value to b:

# define approximation of fit
@jit
def fit(fff, x):
    value = vmap(partial(jnp.interp,xp = xxx))(x = x[:,0], fp = fff)
    return value

b = fit(fff, x)

Then define the same function without @jit, and assign the fitted value to a:

# define approximation of fit
def fit(fff, x):
    value = vmap(partial(jnp.interp,xp = xxx))(x = x[:,0], fp = fff)
    return value

a = fit(fff, x)

Then print out the value of a and b:

a 
DeviceArray([-3.708324e+29], dtype=float32)

b
DeviceArray([-3.7083256e+29], dtype=float32)

a == b 
DeviceArray([False], dtype=bool)

[jakevdp]

Thanks - in the current version of JAX, executed on CPU, I find that the results are equal (note the different import for partial, because it has been removed from the jax namespace in recent versions):

from functools import partial
from jax import vmap, jit
import jax.numpy as jnp
xxx = jnp.array([  0.      ,   0.831025,   3.3241  ,   7.479224,  13.2964  ,
              20.775625,  29.916897,  40.720222,  53.1856  ,  67.31303 ,
              83.1025  , 100.55403 , 119.66759 , 140.44324 , 162.88089 ,
             186.98064 , 212.7424  , 240.16623 , 269.2521  , 300.00003 ])

fff = jnp.array([[ 3.61626554e+00,  3.61627531e+00,            -jnp.inf,
                   -jnp.inf,            -jnp.inf,            -jnp.inf,
        -7.35202107e+37, -7.34325502e+37, -3.50007407e+37,
        -1.64472493e+37, -8.36250276e+36, -2.57747340e+36,
        -5.69342370e+35, -8.81759238e+34, -1.08298795e+34,
        -1.32660759e+32, -2.15171236e+30,  3.61739659e+00,
         3.61751652e+00,  3.61758924e+00]])

x = jnp.array([[2.3543993e+02, 2.0000000e+01, 0.0000000e+00, 0.0000000e+00,
              1.0000000e+00, 1.0000000e+00, 1.1188778e-02]])

# define approximation of fit
def fit(fff, x):
    value = vmap(partial(jnp.interp,xp = xxx))(x = x[:,0], fp = fff)
    return value

a = fit(fff, x)
b = jit(fit)(fff, x)

print(a == b)
# [ True]

But looking at your outputs, it seems like they match to within a reasonable precision for 32-bit floating point arithmetic. It is generally not expected that JIT and non-JIT codepaths will give bitwise-equivalent results, due to the fact that XLA compilation may re-order floating point operations.