Speed up one step NUTS within custome sampler

Bsically, I'd like to use NUTS to sample my 1 of my parameters. Here, I use the tensorflow implementation, but the problem is that it is too painful to run even one iteration. Could anyone help me speed it up?

Some pseudo code:

for gg in range(1,iter_num):
    # 1. sample parameters except "r" by my own sampler

    # 2. sample parameter "r" by NUTS
    # very SLOW...

    target_lpdf = lambda r: log_pdf(r, other_param_prev)
    nuts = tfp.mcmc.NoUTurnSampler(target_log_prob_fn = target_lpdf,
                                       step_size = tf.cast(.1, tf.float32))
    state = tfp.mcmc.sample_chain(num_results=1,
                                      current_state = tf.cast(r_trace[gg-1], tf.float32),
                                      kernel=nuts)
    r_trace[gg] = np.array(state.all_states, dtype = 'float64')

Start with the custom sampler loop in the issue's pseudo-code, especially construction of NoUTurnSampler and the single-result sample_chain call. Run a minimal one-iteration benchmark with the supplied state and target_log_prob_fn, then compare the timing of sampler construction and execution. Done means the runtime bottleneck is isolated and a concrete performance change is identified.