Variational inference with Truncated Normal distribution fails when jit_compile=True

The code below can be successfully executed if I remove the jit_compile=True option from the @tf.function() decorator or if I use a LogNormal distribution instead of a Truncated Normal distribution for my surrogate posterior. If I use Truncated Normal distribution with jit_compile=True I am getting the error:

TypeError: Expected int for argument 'seed2' not <tf.Tensor 'TruncatedNormal/sample/parameterized_truncated_normal/mod:0' shape=(2,) dtype=int32>.

import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions
tfb = tfp.bijectors
dtype = tf.float32


def build_model_log_prob_fn(lam_0, sig, ts, ys):
    model = tfd.JointDistributionSequentialAutoBatched([
        tfd.Exponential(rate=lam_0, name='w', force_probs_to_zero_outside_support=True),
        lambda w: tfd.Normal(loc=w * ts, scale=sig, name='ys')
    ])

    model_log_prob_fn = lambda *x: model.log_prob(x + (ys,))

    return model, model_log_prob_fn


def build_surrogate_posterior(type_: str = 'truncated_normal'):
    """ Build surrogate posterior with trainable parameters """
    
    assert type_ in ('truncated_normal', 'log_normal')
    
    if type_ == 'log_normal':
        Q = tfd.JointDistributionSequentialAutoBatched([
            tfd.LogNormal(
                loc=tf.Variable(tf.constant([0], dtype), name='q_z_loc'),
                scale=tfp.util.TransformedVariable(tf.constant([1], dtype), tfb.Softplus(), name='q_z_scale'))
        ])

    else:
        Q = tfd.JointDistributionSequentialAutoBatched([
            tfd.TruncatedNormal(
                loc=tf.Variable(tf.constant([0], dtype), name='q_z_loc'),
                scale=tfp.util.TransformedVariable(tf.constant([1], dtype), tfb.Softplus(), name='q_z_scale'),
                low=0.,
                high=20.)
        ])

    return Q


lam_0 = tf.constant(2, dtype=dtype)
sig = tf.constant(3, dtype=dtype)
ts = tf.constant([33., 56., 36.], dtype=dtype)
ys = tf.constant([18., 26., 23.], dtype=dtype)

q = build_surrogate_posterior('truncated_normal')   # 'log_normal', 'truncated_normal'

model, target_log_prob_fn = build_model_log_prob_fn(lam_0, sig, ts, ys)

optimizer = tf.optimizers.Adam(learning_rate=1e-2)

@tf.function(jit_compile=True)
def fit_vi():
    return tfp.vi.fit_surrogate_posterior(
        target_log_prob_fn=target_log_prob_fn,
        surrogate_posterior=q,
        optimizer=optimizer,
        num_steps=5000
    )

losses = fit_vi()

Start by running the supplied Python reproducer with TruncatedNormal, fit_vi, and @tf.function(jit_compile=True`, then compare the failing sampling path with the working LogNormal case. Trace how TruncatedNormal produces the seed2 argument under XLA; done means the example runs with jit_compile=True without the TypeError and the regression is covered.