Trying to compile HMC with Tensor arguments - Dynamic inferencing name: "pad" is not supported

It would be desirable to run multiple times the same HMC procedure, but with different arguments. For example one might want to compile the graph (by passing one iteration as an argument) before running it.

Following https://pgaleone.eu/tensorflow/tf.function/2019/04/03/dissecting-tf-function-part-2/#using-a-python-native-type it seems that if the developer passes number of steps as an integer, each time he changes the number, the graph will be recompiled.

Therefore I wanted to pass number of steps as a Tensor. The compilation fails:

InvalidArgumentError: Dynamic inferencing name: "pad"
opcode: "pad"
shape {
  element_type: S32
  dimensions: 2
  layout {
    minor_to_major: 0
    format: DENSE
  }
  is_dynamic_dimension: false
}
metadata {
  op_type: "Pad"
  op_name: "mcmc_sample_chain/trace_scan/size0/Pad"
}
padding_config {
  dimensions {
    edge_padding_high: 1
  }
}
id: 257
operand_ids: 255
operand_ids: 256
frontend_attributes {
}
 is not supported
	while evaluating input dynamism1 of TensorListReserve
	 [[{{node mcmc_sample_chain/trace_scan/TensorArrayV2}}]] [Op:__inference_function_1617]

The example works well without the compilation, and if you compile with the native Python types. How to make it work?

Minimal reproducible example:

N_CHAINS = 4

import tensorflow_probability as tfp
import tensorflow as tf
from functools import partial

def run_nuts_template(
    trace_fn,
    target_log_prob_fn,
    inits,
    bijectors_list,
    num_steps,
    num_burnin,
    step_size=None
):

    kernel = tfp.mcmc.NoUTurnSampler(
        target_log_prob_fn,
        step_size=step_size
    )
    res = tfp.mcmc.sample_chain(
        num_results=num_steps,
        num_burnin_steps=num_burnin,
        current_state=inits,
        kernel=kernel,
        trace_fn=trace_fn
    )
    return res

def joint_log_prob(
        sigma_obs,
        product_obs
):
    tdn = tfp.distributions.Normal(10,10)
    so = tdn.log_prob(sigma_obs)
    po = tdn.log_prob(product_obs)
    return tf.reshape((so + po), (-1,))


def trace_fn(_, pkr):
    return (
        pkr.target_log_prob,
        pkr.leapfrogs_taken,
        pkr.has_divergence,
        pkr.energy,
        pkr.log_accept_ratio,
        pkr.is_accepted,
    )


def get_inits_and_step_sizes():

    initial_chain_state = [
        tf.ones([N_CHAINS, 1, 1], dtype=tf.float32, name="init_sigma_obs"),
        tf.ones([N_CHAINS, 1, 1], dtype=tf.float32, name="init_product_obs")
    ]

    step_size = [
        tf.ones([N_CHAINS, 1, 1], dtype=tf.float32),
        tf.ones([N_CHAINS, 1, 1], dtype=tf.float32)
    ]

    return initial_chain_state, step_size

initial_chain_state, step_sizes = get_inits_and_step_sizes()
unconstraining_bijectors = [
    tfp.bijectors.Softplus(),
    tfp.bijectors.Softplus()
]


run_nuts = partial(run_nuts_template, trace_fn)
run_nuts_opt = tf.function(
    run_nuts,
    autograph=True,
    experimental_compile=True
)

res = run_nuts_opt(
    joint_log_prob,
    initial_chain_state,
    unconstraining_bijectors,
    num_steps=tf.constant(20, dtype=tf.int32),
    num_burnin=tf.constant(20, dtype=tf.int32),
    step_size=step_sizes
)

Versions

Tensorflow 2.4.0
Tensorflow Probability 0.10.1

Ran on Ubuntu 18.04

Start with the minimal reproducible example at run_nuts_opt and tfp.mcmc.sample_chain using TensorFlow 2.4.0 and TensorFlow Probability 0.10.1, then trace the TensorListReserve/Pad compilation failure. Done means the same HMC procedure compiles and runs when num_steps and num_burnin are Tensor arguments, while retaining the behavior that works with native Python integers.