Get jit running on Gaussian Process example

numpyro
#1

Hi,

I try to get JIT running for the Gaussian Process example but the kernel code
make a trouble due to the if condition:

    if include_noise:
        k += (noise + jitter) * jnp.eye(X.shape[0])

does anyone have tried?
Thanks

#2

For static arguments, you can specify them in the static_argnums argument of jit. See its docs.

#3

Got it , thanks.