def __call__(self, shape, dtype=None):
if not self.built:
self.build(shape, dtype)
return generated_random_variables.Independent(
generated_random_variables.HalfCauchy(
loc=self.loc, scale=self.scale).distribution,
reinterpreted_batch_ndims=len(shape))
def __call__(self, x):
"""Computes regularization using an unbiased Monte Carlo estimate."""
prior = generated_random_variables.Independent(
generated_random_variables.HalfCauchy(
loc=tf.broadcast_to(self.loc, x.distribution.event_shape),
scale=tf.broadcast_to(self.scale, x.distribution.event_shape)
).distribution,
reinterpreted_batch_ndims=len(x.distribution.event_shape))
negative_entropy = x.distribution.log_prob(x)
cross_entropy = -prior.distribution.log_prob(x)
return self.scale_factor * (negative_entropy + cross_entropy)