sample_annealed_importance_chain example consistently gives bad results

Hi! I'm trying to use sample_annealed_importance_chain. I tried the example and modified it slightly to get it actually running (add a comma and replace tf.log/lgamma with tf.math.log/lgamma):

tfd = tfp.distributions

# Run 100 AIS chains in parallel
num_chains = 32
dims = 20
dtype = np.float32

proposal = tfd.MultivariateNormalDiag(
   loc=tf.zeros([dims], dtype=dtype))

target = tfd.TransformedDistribution(
  distribution=tfd.Sample(
      tfd.Gamma(concentration=dtype(2), rate=dtype(3)),
      sample_shape=[dims]),
  bijector=tfp.bijectors.Invert(tfp.bijectors.Exp()))

chains_state, ais_weights, kernels_results = (
    tfp.mcmc.sample_annealed_importance_chain(
        num_steps=1000,
        proposal_log_prob_fn=proposal.log_prob,
        target_log_prob_fn=target.log_prob,
        current_state=proposal.sample(num_chains),
        make_kernel_fn=lambda tlp_fn: tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=tlp_fn,
          step_size=0.2,
          num_leapfrog_steps=2)))

log_estimated_normalizer = (tf.reduce_logsumexp(ais_weights)

                            - np.log(num_chains))
log_true_normalizer = tf.math.lgamma(2.) - 2. * tf.math.log(3.)
print(tf.math.exp(log_estimated_normalizer), tf.math.exp(log_true_normalizer))

Unfortunately, the estimated normalizing constant is consistently one order of magnitude higher then it should be and the number of steps doesn't really seem to matter, I've tried different values from 1 to 1000. I would expect the variance of the estimate to get lower as num_steps is decreased, as described in the documentation.
Also, the note in the documentation says that the distribution functions are only called three times in graph mode. How is that possible? I would expect them to be called at every HMC step.

Start by running the provided sample_annealed_importance_chain example and compare its estimated normalizing constant with the analytical value. Then inspect the sample_annealed_importance_chain entry point and its HMC execution path, including the graph-mode behavior. Done means the example produces a correct estimate and the documentation explains the number of distribution-function calls.