Enforce Distribution statistics' batch shape contract by broadcasting.

The contract is that a Distribution object is supposed to look like
all of its parameters' batch shapes are broadcasted together to form
the batch shape of the Distribution.

This mostly happens automatically, but some statistics don't actually
depend on all of the respective Distribution's parameters.  (The most
common case is the variance of a location-scale family doesn't depend
on the location.)  This CL inserts tf.broadcast_to at the outputs of
appropriate statistic methods to enforce the contract, and updates a
Hypothesis test to check for shape-correctness of instantiable
Distributions mechanically.

The Distribution methods affected are
- GeneralizedPareto.entropy
- PERT.mode
- Triangular.entropy
- LambertWNormal.{variance,stddev}
- ExponentiallyModifiedGaussian.{mean,variance,stddev}

PiperOrigin-RevId: 383723788

Author: axch

tensorflow_probability/python/distributions/distribution_properties_test.py

@@ -50,11 +50,6 @@
 
 
 STATISTIC_CONSISTENT_SHAPES_TEST_BLOCK_LIST = (
-    'BatchReshape',  # b/183405889
-    'Independent',  # b/183405889
-    'Mixture',  # b/183405889
-    'Sample',  # b/183405889
-    'TransformedDistribution',  # b/183405889
 )
 
 
@@ -139,13 +134,12 @@ def check_statistic(
     try:
       with tfp_hps.no_tf_rank_errors():
         result = getattr(dist, statistic)()
-      msg = 'Shape {} not compatible with expected {}.'.format(
-          result.shape, expected_static_shape)
-      self.assertTrue(expected_static_shape.is_compatible_with(
-          tf.broadcast_static_shape(result.shape, expected_static_shape)), msg)
+      msg = 'Shape {} of {} not compatible with expected {}.'.format(
+          result.shape, statistic, expected_static_shape)
+      self.assertTrue(
+          expected_static_shape.is_compatible_with(result.shape), msg)
       self.assertAllEqual(self.evaluate(expected_dynamic_shape),
-                          self.evaluate(tf.broadcast_dynamic_shape(
-                              tf.shape(result), expected_dynamic_shape)))
+                          self.evaluate(tf.shape(result)))
     except NotImplementedError:
       pass
 

tensorflow_probability/python/distributions/exponentially_modified_gaussian.py

@@ -198,10 +198,13 @@ def _log_cdf(self, x):
         -u + vsquared / 2. + special_math.log_ndtr((u - vsquared) / v))
 
   def _mean(self):
-    return self.loc + 1 / self.rate
+    return tf.broadcast_to(
+        self.loc + 1 / self.rate, self._batch_shape_tensor())
 
   def _variance(self):
-    return tf.square(self.scale) + 1 / tf.square(self.rate)
+    return tf.broadcast_to(
+        tf.square(self.scale) + 1 / tf.square(self.rate),
+        self._batch_shape_tensor())
 
   def _parameter_control_dependencies(self, is_init):
     assertions = []

tensorflow_probability/python/distributions/generalized_pareto.py

@@ -290,7 +290,8 @@ def result():
       return result()
 
   def _entropy(self):
-    return tf.math.log(self.scale) + self.concentration + 1
+    ans = tf.math.log(self.scale) + self.concentration + 1
+    return tf.broadcast_to(ans, self._batch_shape_tensor())
 
   # TODO(b/145620027): Finalize choice of bijector.
   def _default_event_space_bijector(self):

tensorflow_probability/python/distributions/lambertw_f.py

@@ -269,17 +269,18 @@ def _variance(self):
         tf.convert_to_tensor(np.inf, dtype=self.dtype))
 
     if self.allow_nan_stats:
-      return tf.where(
+      ans = tf.where(
           tailweight < 1.0,
           result_where_defined,
           tf.convert_to_tensor(np.nan, self.dtype))
     else:
-      return distribution_util.with_dependencies([
+      ans = distribution_util.with_dependencies([
           assert_util.assert_greater_equal(
               tf.ones([], dtype=self.dtype),
               tailweight,
               message="variance not defined for components of tailweight >= 1"),
       ], result_where_defined)
+    return tf.broadcast_to(ans, self._batch_shape_tensor())
 
   def _mode(self):
     # Mode always exists (for any tail parameter) and equals the location / mean

tensorflow_probability/python/distributions/pert.py

@@ -213,7 +213,8 @@ def _mean(self):
     return self._transformed_beta().mean()
 
   def _mode(self):
-    return tf.convert_to_tensor(self.peak)
+    return tf.broadcast_to(tf.convert_to_tensor(self.peak),
+                           self._batch_shape_tensor())
 
   def _variance(self):
     low = tf.convert_to_tensor(self.low)

tensorflow_probability/python/distributions/triangular.py

@@ -251,7 +251,8 @@ def _cdf(self, x):
     return tf.where(x >= high, tf.ones_like(x), result_if_not_big)
 
   def _entropy(self):
-    return 0.5 - np.log(2.) + tf.math.log(self.high - self.low)
+    ans = 0.5 - np.log(2.) + tf.math.log(self.high - self.low)
+    return tf.broadcast_to(ans, self._batch_shape_tensor())
 
   def _mean(self):
     return (self.low + self.high + self.peak) / 3.