A pre-canned builder for a ubiquitous stochastic KL estimator.

def stochastic_kl_builder(posterior, prior, sample):
  """A pre-canned builder for a ubiquitous stochastic KL estimator."""
  return tf.subtract(
      tf.reduce_sum(posterior.log_prob(sample)),
      tf.reduce_sum(prior.log_prob(sample)))