HMC sampling results differ when using / not using XLA compilation

In the following code, I get an acceptance rate of one when decorating run_chain just with @tf.function, but when enabling XLA compilation by using @tf.function(experimental_compile=True), the acceptance rate becomes zero.
I stumbled upon this with TFP 0.11.1 and TF 2.3.1 on Python 3.6.9 on a Ubuntu machine, but also the nightly builds seem to show that behavior, as can be checked using this Colab notebook.

In case you think the log-probability is weird: this issue popped up in a larger model, and this is a minimal-ish version which reproduces the problem. Rest assured, the complete model makes sense ;-)

Any insight into this issue would be greatly appreciated!

import numpy as np

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
import tensorflow_probability as tfp

tfd = tfp.distributions

n = 23

mask = np.ones((n,n))
mask[np.diag_indices(n)] = 0.0

def log_prob(X):
    d0 = 2.0
    # d: pairwise distance matrix
    d = tf.linalg.norm(tf.expand_dims(X, 0) - tf.expand_dims(X, 1), axis=2)
    bla = tf.math.square((d0 - d) * mask)
    return -tf.reduce_sum(bla)
                   
initial_state = np.random.uniform(low=-10, high=10,
                                  size=(n, 3)).astype(np.float64)

hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=log_prob,
    step_size=.005,
    num_leapfrog_steps=10
    )


@tf.function(experimental_compile=True)
def run_chain(initial_state, num_results=1000, num_burnin_steps=100):
    return tfp.mcmc.sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=initial_state,
        kernel=hmc_kernel,
        trace_fn=lambda current_state, kernel_results: kernel_results
        )


samples, kernel_results = run_chain(initial_state, num_results=1000)

print(kernel_results.is_accepted.numpy().mean())

Run the supplied Python reproducer, comparing run_chain with @tf.function and experimental_compile=True; the linked Colab notebook provides the same setup. Inspect the HMC kernel results and XLA execution path, with the work complete when the acceptance-rate discrepancy has a confirmed cause and a reproducible fix or documented limitation.