Add entropy estimate to the tanh transformed distribution.
@@ -190,6 +190,16 @@ class TanhTransformedDistribution(tfd.TransformedDistribution):
-inverse_threshold) - log_epsilon
self._log_prob_right = self.distribution.log_survival_function(
inverse_threshold) - log_epsilon
+ self._entropy_estimate = self._compute_entropy_estimate()
+
+ def _compute_entropy_estimate(self):
+ """Computes the estimate of the distribution entropy using three points."""
+ mean = self.distribution.mean()
+ std = self.distribution.stddev()
+ estimation_points = jnp.array([mean - std, mean, mean + std])
+ return (self.distribution.entropy() + jnp.mean(
+ self.bijector.forward_log_det_jacobian(
+ estimation_points, event_ndims=0)))
def log_prob(self, event):
# Without this clip there would be NaNs in the inner tf.where and that
@@ -205,6 +215,9 @@ class TanhTransformedDistribution(tfd.TransformedDistribution):
def mode(self):
return self.bijector.forward(self.distribution.mode())
+ def entropy(self):
+ return self._entropy_estimate
+
@classmethod
def _parameter_properties(cls, dtype: Optional[Any], num_classes=None):
td_properties = super()._parameter_properties(dtype,
