KLDivergence on output layer from trainable prior causes freezing

The following contrived snippet will freeze after some random number of epochs without displaying any error messages. I came across this bug while testing something; it's not blocking me. I just want to share one way to break tfp/tf without generating any errors ;)

loc = tf.Variable(tf.random.normal([ndim], stddev=0.1, dtype=tf.float32))
scale = tfp.util.TransformedVariable(
    tf.random.normal([ndim], mean=1.0, stddev=0.1, dtype=tf.float32),
    bijector=tfb.Chain([tfb.Shift(1e-5), tfb.Softplus(), tfb.Shift(0.5413)]))
prior = tfd.Independent(tfd.Normal(loc=loc, scale=scale), reinterpreted_batch_ndims=1)


_input = tfkl.Input(shape=(1,))
_loc = tfkl.Dense(ndim, name="loc_params")(_input)
_scale = tfkl.Dense(ndim, name="untransformed_scale_params")(_input)
_scale = tf.math.softplus(_scale + np.log(np.exp(1) - 1)) + 1e-5
_output = tfpl.DistributionLambda(
    make_distribution_fn=lambda t: tfd.Independent(tfd.Normal(loc=t[0], scale=t[1])),
    activity_regularizer=tfpl.KLDivergenceRegularizer(prior, use_exact_kl=True, weight=0.1)
)([_loc, _scale])
model = tf.keras.Model(_input, _output)
model.compile(optimizer='adam', loss=lambda y_true, model_out: -model_out.log_prob(y_true))
hist = model.fit(ds, epochs=N_EPOCHS, verbose=2)

If I set N_EPOCHS to something really small (e.g., 5) then it's likely training will complete and I can see that the trainable prior has updated its loc and scale, so the mechanism is working.

It seems to be specific to having a trainable prior.

Here is a runnable gist. You can change PRIOR_TRAINABLE = False and see that it completes.

It also seems to be specific when KL divergence is used in the output layer used for the log-likelihood calculation. In this notebook (Colab link), the prior can be made trainable but KL divergence is calculated only on a latent bottleneck, and training completes.

Run the linked gist with PRIOR_TRAINABLE enabled, then compare it with PRIOR_TRAINABLE=False and the linked betaVAE notebook where KL divergence is applied at a latent bottleneck. Start at tfpl.KLDivergenceRegularizer used by tfpl.DistributionLambda and trace the output-layer training path. Done means the provided training example completes reliably with a trainable prior and the regression is covered by a test.