Lots of rejections samplers... any ideas?

[SamDuffield]

I am trying to implement a fairly involved method to sample from a joint distribution based on chained rejection samplers .

In theory there should be a significant performance increase for a large number of samples over the much simpler full resampling method, however in current JAX implementation below I am not seeing this.

The end product (I've posted a simpler version below) actually runs this method over many dimensions and gets placed within a higher scan, this is where I am seeing particularly poor performance.

Some comments on the implementation:

• Some of the rejection samplers will accept much earlier than others, to deal with this (without using non-concrete indexing etc) I run while_loop over map over cond where the negative execution of the cond is blank but the overheads here may be costly?
• I use map rather than vmap to ensure the positive cond isn't also executed
• The bound for the rejection sampling is found adaptively - so the rejection samplers need to communicate to check if we need to restart
• I am not sure if I should use jit here somewhere to avoid recompilation? Particularly as it's going inside scan afterwards

I realise this sort of algorithm isn't particularly well suited to JAX's control flow, but just thought I'd ask if anyone had any ideas for performance enhancements?

from time import time
from jax import numpy as np, random, vmap
from jax.lax import while_loop, cond, map

initial_keys = random.split(random.PRNGKey(0), 3)

n = 10000

x0_samps = random.normal(initial_keys[0], shape=(n,))
x1_samps = random.normal(initial_keys[1], shape=(n,))

rho = 0.9


def log_conditional_density(x0_given, x1):
    return -0.5 / (1 - rho ** 2) * (x1 - rho * x0_given) ** 2


def full_resample_single(x0_single, x1_all, random_key):
    log_weights = vmap(log_conditional_density, (None, 0))(x0_single, x1_all)
    return x1_all[random.categorical(random_key, log_weights)]


def full_resampling(x0_all, x1_all, random_key):
    return vmap(full_resample_single, (0, None, 0))(x0_all, x1_all, random.split(random_key, len(x0_all)))


def full_resample_single_cond(not_yet_accepted, x0_single, x1_all, random_key):
    return cond(not_yet_accepted,
                lambda tup: full_resample_single(*tup),
                lambda _: x1_all[0],  # just dummy of same type as full_resample_single output
                (x0_single, x1_all, random_key))


def rejection_proposal_single(x0_single, x1_all, bound, random_key):
    random_key, choice_key, uniform_key = random.split(random_key, 3)
    x1_single = x1_all[random.choice(choice_key, len(x1_all))]
    conditional_dens = np.exp(log_conditional_density(x0_single, x1_single))
    return x1_single, conditional_dens, random.uniform(uniform_key) > conditional_dens / bound, random_key


def rejection_proposal_single_cond(not_yet_accepted, x0_single, x1_all, bound, random_key):
    return cond(not_yet_accepted,
                lambda tup: rejection_proposal_single(*tup),
                lambda _: (x1_all[0], 0., False, random_key), # just dummy of same type as rejection_proposal_single output
                (x0_single, x1_all, bound, random_key))


def rejection_proposal_all(not_yet_accepted_arr, x1_all_sampled, x0_all, x1_all, bound, random_keys, rejection_iter):
    n = len(x0_all)
    mapped_tup = map(lambda i: rejection_proposal_single_cond(not_yet_accepted_arr[i],
                                                              x0_all[i],
                                                              x1_all,
                                                              bound,
                                                              random_keys[i]), np.arange(n))
    x1_all_proposed, dens_evals, not_yet_accepted_arr_new, random_keys = mapped_tup

    x1_all_sampled = np.where(not_yet_accepted_arr, x1_all_proposed, x1_all_sampled)

    # Check if we need to start again
    max_dens = np.max(dens_evals)
    reset_bound = max_dens > bound
    bound = np.where(reset_bound, max_dens, bound)
    not_yet_accepted_arr_new = np.where(reset_bound, np.ones(n, dtype='bool'), not_yet_accepted_arr_new)
    return not_yet_accepted_arr_new, x1_all_sampled, x0_all, x1_all, bound, random_keys, rejection_iter + 1


def rejection_resampling(x0_all, x1_all, random_key, maximum_rejections):
    init_key, while_key = random.split(random_key, 2)
    n = len(x0_all)

    # Prerun to initiate bound
    x1_initial = x1_all
    initial_cond_dens = np.exp(vmap(log_conditional_density)(x0_all, x1_initial))
    initial_bound = np.max(initial_cond_dens)
    initial_not_yet_accepted_arr = random.uniform(init_key, (n,)) > initial_cond_dens / initial_bound

    out_tup = while_loop(lambda tup: np.logical_and(tup[0].sum() > 0, tup[6] < maximum_rejections),
                         lambda tup: rejection_proposal_all(*tup),
                         (initial_not_yet_accepted_arr,
                          x1_initial,
                          x0_all,
                          x1_all,
                          initial_bound,
                          random.split(while_key, n),
                          1))
    not_yet_accepted_arr, final_rej_particles, _, _, final_bound, random_keys, rej_attemped = out_tup

    full_particles = map(lambda i: full_resample_single_cond(not_yet_accepted_arr[i],
                                                             x0_all[i],
                                                             x1_all,
                                                             random_keys[i]), np.arange(n))
    merged_particles = np.where(not_yet_accepted_arr, full_particles, final_rej_particles)

    return merged_particles, not_yet_accepted_arr, final_bound, rej_attemped


fr_key, rr_key = random.split(initial_keys[2])

fr_start = time()
full_resamples = full_resampling(x0_samps, x1_samps, fr_key).block_until_ready()
fr_end = time()
fr_time = fr_end - fr_start
print(fr_time)

rr_start = time()
rejection_resamples, final_not_yet_accepted, final_bound, rejections_attemped = rejection_resampling(x0_samps, x1_samps,
                                                                                                     rr_key, 100)
rejection_resamples.block_until_ready()
rr_end = time()
rr_time = rr_end - rr_start
print(rr_time)

print(final_not_yet_accepted.sum())