SimpleStepSizeAdaptation HMC leads to NaN values for AR1 parameter estimation

[Qiuliang Ye]

Hi everyone, recently I used the tfp.mcmc.SimpleStepSizeAdaptation for parameter estimation in AR1 and found that NaN values happen, while the tfp.mcmc.RandomWalkMetropolis does not show this problem.

I implemented 10 independent chains with random initialization states:
tf.concat([tf.random.uniform([N_CHAINS, 1]), # for sigma^2 tf.random.normal([N_CHAINS, 1]), # for delta tf.random.uniform([N_CHAINS, 1])-0.5], axis=-1) # for AR1 coefficient

The adaptation is:
TransformedTransitionKernel(HamiltonianMonteCarlo(posterior_dist(replicate_observations), step_size=0.01, num_leapfrog_steps=3)

I replicate the problem in a Colab: https://colab.research.google.com/drive/1i0AATZFozKVi6Lq5jYKV6C0kMWwHIE_7?usp=drive_link

I checked that the parameters in some chains go to infinity at the very beginning stage, and I dont know why. Could anyone help me?