tfp.experimental.nn.losses.compute_extra_loss( layer, loss_fn=tfp.experimental.nn.losses.kl_divergence_monte_carlo )