Does this code is inefficient for JIT?

[jecampagne]

Dear experts,

Here is a simple JAX code to show the Metropolis Algorithm in action to solve a 3 parameters bayesian regression pb. Running wo the JIT compilation is ok even on a CPU. Now I would like to know why when the 2 line concerning JIT are decommented then the timing is not really different both of CPU (Jit or not JIT) and comparing running on a CPU or on a K80/Nvidia GPU?

May I have coded in a wrong/inefficient way?

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import jacfwd, jacrev, hessian
from jax.ops import index, index_update
from functools import partial

import scipy.stats as scs
import numpy as np

#@partial(jax.jit, static_argnums=(1,))
def jax_metropolis_kernel(rng_key, logpdf, position, log_prob):
    key, subkey = jax.random.split(rng_key)
    """Moves the chain by one step using the Random Walk Metropolis algorithm.
    """
    move_proposals = jax.random.normal(key, shape=position.shape) * 0.1
        
    proposal = position + move_proposals
    proposal_log_prob = logpdf(proposal)

    log_uniform = jnp.log(jax.random.uniform(subkey))
    do_accept = log_uniform < proposal_log_prob - log_prob

    position = jnp.where(do_accept, proposal, position)
    log_prob = jnp.where(do_accept, proposal_log_prob, log_prob)
    return position, log_prob

#@partial(jax.jit, static_argnums=(1, 2))
def jax_metropolis_sampler(rng_key, n_samples, logpdf, initial_position):
    """Generate samples using the Random Walk Metropolis algorithm.
    """
    def mh_update(i, state):
        key, positions, log_prob = state
        _, key = jax.random.split(key)
                
        new_position, new_log_prob = jax_metropolis_kernel(key, 
                                                           logpdf, 
                                                           positions[i-1], 
                                                           log_prob)
                
        
        positions=positions.at[i].set(new_position)
        return (key, positions, new_log_prob)

    #Initialisation
    keys = jax.random.split(rng_key,num=4)
    all_positions = jnp.zeros((n_samples,initial_position.shape[0]))  # 1 chain for each vmap call    ?
    all_positions=all_positions.at[0,0].set(jax.random.normal(keys[0])+1.)
    all_positions=all_positions.at[0,1].set(jax.random.normal(keys[1])+2.)
    all_positions=all_positions.at[0,2].set(jax.random.uniform(keys[2],minval=1.0, maxval=3.0))


    logp = logpdf(all_positions[0])
    
    initial_state = (rng_key,all_positions, logp)
    rng_key, all_positions, log_prob = jax.lax.fori_loop(1, n_samples, 
                                                 mh_update, 
                                                 initial_state)
    
    return all_positions

def jax_my_logpdf(par,xi,yi):
    # priors: a=par[0], b=par[1], sigma=par[2]
    logpdf_a = jax.scipy.stats.norm.logpdf(x=par[0],loc=1.,scale=1.)
    logpdf_b = jax.scipy.stats.norm.logpdf(x=par[1],loc=2.,scale=1.)
    logpdf_s = jax.scipy.stats.gamma.logpdf(x=par[2],a=3,scale=1.)

    val = xi*par[1]+par[0]
    tmp = jax.scipy.stats.norm.logpdf(x=val,loc=yi,scale=par[2])    
    log_likeh= jnp.sum(tmp)
    
    rc = log_likeh + logpdf_a + logpdf_b + logpdf_s

    return log_likeh + logpdf_a + logpdf_b + logpdf_s

######## Main ########
n_dim = 3
n_forget = 1_000
n_samples = 100_000 + n_forget
n_chains = 100
rng_key = jax.random.PRNGKey(42)

# the data xi, yi
sample_size = 5_000
sigma_e = 1.5             # true value of parameter error sigma
random_num_generator = np.random.RandomState(0)
xi = 10.0 * random_num_generator.rand(sample_size)
e = random_num_generator.normal(0, sigma_e, sample_size)
yi = 1.0 + 2.0 * xi +  e          # a = 1.0; b = 2.0; y = a + b*x


rng_keys = jax.random.split(rng_key, n_chains)    # generate an array of size (n_chains, 2)
initial_position = jnp.ones((n_dim, n_chains))    # generate an array of size (n_dim, n_chains)
                                                  # so for vmap one should connect axis 0 of rng_keys  
                                                  # and axis 1 of initial_position

#print("main initial_position shape",initial_position.shape)

run_mcmc = jax.vmap(jax_metropolis_sampler, 
                    in_axes=(0, None, None, 1),   # see comment above 
                    out_axes=0)                   # output axis 0 hold the vectorization over n_chains
                                                  # => (n_chains, n_samples, n_dims)


all_positions = run_mcmc(rng_keys, n_samples, 
                     lambda par: jax_my_logpdf(par,xi,yi), 
                     initial_position)

then, after a first call to get rid of compilation time

%timeit all_positions = run_mcmc(rng_keys, n_samples, 
                     lambda par: jax_my_logpdf(par,xi,yi), 
                     initial_position)

I get on CPU: 1min 27s (mean of 7 runs) without JIT and 1min 20s with JIT.

Thanks for your advises.

[jakevdp]

How are you doing the timing? Are you using block_until_ready() to account for Asynchronous Dispatch? Are you including JIT compilation time in your timing of the JIT version? More information here would help us provide better assistance.

[jecampagne]

Well, the timing that I have done is to count the wall time (on my watch) to get the all_positions array once I have activates the jupyter cell containing the "main" part. For sure it is not a very precise answer and probably you have better way to get this wall time.

[jecampagne]

Sorry, I have updated the code for the (xi,yi) data set.

[jakevdp]

OK, I looked at the code - I'm not sure about performance yet, but I notice that you're using scipy.stats within a vmapped function, and I suspect that's not doing what you want. As currently written, every chain has the same starting value (because scipy.stats functions are evaluated statically). If you want random start locations for each chain, you should use jax.random with a different key for each vmap entry.

Just a quick demonstration of what I'm talking about:

def f(x):
  return scs.norm.rvs(loc=1,scale=1)

vmap(f)(jnp.arange(5))
# DeviceArray([1.6125425, 1.6125425, 1.6125425, 1.6125425, 1.6125425], dtype=float32)

[jecampagne]

I have updated the code according to your kind advises: I have switched from scipy.stats to jax.random code for initialisation and also give the timing results on CPU by using "%timeit" with and without JIT .

[jakevdp]

Thanks - looking more closely now, it seems that the bulk of your computation is done within the function you pass to fori_loop, and this function will be JIT-compiled by default. Because of that, adding the outer jit wrapper is not going to make much difference: the important pieces of your code are already jitted.

[jecampagne]

Hum, ok I see. Thanks for your clear explanation.