Avoid creating new computation nodes in `kernel_results` from MCMC kernels

The one_step method of some MCMC kernels seem to return a kernel_result that contains new computation nodes.

If I run the following loop in graph mode, I would create new computation nodes in every iteration.

1. boot strap and get "kernel result 0"
2. Use "kernel result 0" to run the one_step method to get "kernel result 1"
3. Use "kernel result 1" to run the one_step method to get "kernel result 2"
   ...

How to avoid that?

permalink

      return result_state, NUTSKernelResults(
          next_.target_log_prob, next_.grads_target_log_prob,
          leapfrogs_taken + new_leapfrogs,
          leapfrogs_computed + tf.math.reduce_max(input_tensor=new_leapfrogs))

Start with tensorflow_probability/python/experimental/mcmc/nuts.py at the linked line, then trace how one_step constructs and returns NUTSKernelResults in graph mode. Reproduce the iterative loop and determine what behavior the existing MCMC tests establish; done means repeated one_step calls do not create new computation nodes.