Remove broadcasting_shapes(shape, n=1), which directly returns shape. Adds an error to prevent this case in the future.

Add comment showing how to use broadcast_compatible_shapes instead of broadcasting_shapes. Unfortunately doing so would reveal new failures.

(TODO) Verify that the forward transform adds batch/event dims to match the forward transform of broadcasted-to-full-batch/event.

(TODO) Verify that FLDJ [effectively] broadcasts before reducing.

PiperOrigin-RevId: 375137840

Author: brianwa84

tensorflow_probability/python/distributions/distribution_properties_test.py

@@ -429,17 +429,28 @@ def check_event_space_bijector_constrains(self, dist, data):
     if event_space_bijector is None:
       return
 
+    # Draw a sample shape
+    sample_shape = data.draw(tfp_hps.shapes())
+    inv_event_shape = event_space_bijector.inverse_event_shape(
+        tensorshape_util.concatenate(dist.batch_shape, dist.event_shape))
+
+    # Draw a shape that broadcasts with `[batch_shape, inverse_event_shape]`
+    # where `inverse_event_shape` is the event shape in the bijector's
+    # domain. This is the shape of `y` in R**n, such that
+    # x = event_space_bijector(y) has the event shape of the distribution.
+
+    # TODO(b/174778703): Actually draw broadcast compatible shapes.
+    batch_inv_event_compat_shape = inv_event_shape
+    # batch_inv_event_compat_shape = data.draw(
+    #     tfp_hps.broadcast_compatible_shape(inv_event_shape))
+    # batch_inv_event_compat_shape = tensorshape_util.concatenate(
+    #     (1,) * (len(inv_event_shape) - len(batch_inv_event_compat_shape)),
+    #     batch_inv_event_compat_shape)
+
     total_sample_shape = tensorshape_util.concatenate(
-        # Draw a sample shape
-        data.draw(tfp_hps.shapes()),
-        # Draw a shape that broadcasts with `[batch_shape, inverse_event_shape]`
-        # where `inverse_event_shape` is the event shape in the bijector's
-        # domain. This is the shape of `y` in R**n, such that
-        # x = event_space_bijector(y) has the event shape of the distribution.
-        data.draw(tfp_hps.broadcasting_shapes(
-            event_space_bijector.inverse_event_shape(
-                tensorshape_util.concatenate(
-                    dist.batch_shape, dist.event_shape)), n=1))[0])
+        sample_shape, batch_inv_event_compat_shape)
+    # full_sample_batch_event_shape = tensorshape_util.concatenate(
+    #     sample_shape, inv_event_shape)
 
     y = data.draw(
         tfp_hps.constrained_tensors(
@@ -451,6 +462,14 @@ def check_event_space_bijector_constrains(self, dist, data):
       with tf.control_dependencies(dist._sample_control_dependencies(x)):
         self.evaluate(tf.identity(x))
 
+      # TODO(b/158874412): Verify DoF changing default bijectors.
+      # y_bc = tf.broadcast_to(y, full_sample_batch_event_shape)
+      # x_bc = event_space_bijector(y_bc)
+      # self.assertAllClose(x, x_bc)
+      # fldj = event_space_bijector.forward_log_det_jacobian(y)
+      # fldj_bc = event_space_bijector.forward_log_det_jacobian(y_bc)
+      # self.assertAllClose(fldj, fldj_bc)
+
   @parameterized.named_parameters(
       {'testcase_name': dname, 'dist_name': dname}
       for dname in sorted(list(set(dhps.INSTANTIABLE_BASE_DISTS.keys()) -

tensorflow_probability/python/internal/hypothesis_testlib.py

@@ -593,7 +593,7 @@ def broadcasting_named_shapes(draw, batch_shape, param_names):
   n = len(param_names)
   return dict(
       zip(draw(hps.permutations(param_names)),
-          draw(broadcasting_shapes(batch_shape, n))))
+          draw(broadcasting_shapes(batch_shape, n, no_error_n_eq_1=True))))
 
 
 def _compute_rank_and_fullsize_reqd(draw, target_shape, current_shape, is_last):
@@ -651,7 +651,7 @@ def broadcast_compatible_shape(shape):
 
 
 @hps.composite
-def broadcasting_shapes(draw, target_shape, n):
+def broadcasting_shapes(draw, target_shape, n, no_error_n_eq_1=False):
   """Strategy for drawing a set of `n` shapes that broadcast to `target_shape`.
 
   For each shape we need to choose its rank, and whether or not each axis i is 1
@@ -663,12 +663,17 @@ def broadcasting_shapes(draw, target_shape, n):
     draw: Hypothesis strategy sampler supplied by `@hps.composite`.
     target_shape: The target (fully-defined) batch shape.
     n: Python `int`, the number of shapes to draw.
+    no_error_n_eq_1: If True, don't raise ValueError when n==1.
 
   Returns:
     shapes: A strategy for drawing sequences of `tf.TensorShape` such that the
       set of shapes in each sequence broadcast to `target_shape`. The shapes are
       fully defined.
   """
+  if n == 1 and not no_error_n_eq_1:
+    raise ValueError('`broadcasting_shapes(shp, n=1) is just `shp`. '
+                     'Did you want `broadcast_compatible_shape`? If `n` '
+                     'is stochastic, add arg `no_error_n_eq_1=True`.')
   target_shape = tf.TensorShape(target_shape)
   target_rank = tensorshape_util.rank(target_shape)
   result = []