Fix broadcasting of event and samples in Empirical.

PiperOrigin-RevId: 238975493

tensorflow_probability/python/distributions/empirical.py

@@ -31,9 +31,16 @@
 
 def _broadcast_event_and_samples(event, samples, event_ndims):
   """Broadcasts the event or samples."""
-  samples_shape = tf.shape(input=samples)[:-event_ndims - 1]
+  # This is the shape of self.samples, without the samples axis, i.e. the shape
+  # of the result of a call to dist.sample(). This way we can broadcast it with
+  # event to get a properly-sized event, then add the singleton dim back at
+  # -event_ndims - 1.
+  samples_shape = tf.concat(
+      [tf.shape(input=samples)[:-event_ndims - 1],
+       tf.shape(input=samples)[tf.rank(samples) - event_ndims:]],
+      axis=0)
   event *= tf.ones(samples_shape, dtype=event.dtype)
-  event = tf.expand_dims(event, axis=-event_ndims-1)
+  event = tf.expand_dims(event, axis=-event_ndims - 1)
   samples *= tf.ones_like(event, dtype=samples.dtype)
 
   return event, samples
@@ -259,7 +266,7 @@ def _cdf(self, event):
     cdf = tf.reduce_sum(
         input_tensor=tf.cast(
             tf.reduce_all(
-                input_tensor=samples - event <= 0,
+                input_tensor=samples <= event,
                 axis=tf.range(-self._event_ndims, 0)),
             dtype=tf.int32),
         axis=-1) / self.num_samples

tensorflow_probability/python/distributions/empirical_test.py

@@ -650,6 +650,18 @@ def testVarianceAndStd(self):
     self.assertAllClose(self.evaluate(dist.stddev()),
                         np.sqrt(expected_variance))
 
+  def testLogProbAfterSlice(self):
+    samples = np.random.randn(6, 5, 4)
+    dist = empirical.Empirical(samples=samples, event_ndims=1)
+    self.assertAllEqual((6,), dist.batch_shape)
+    self.assertAllEqual((4,), dist.event_shape)
+    sliced_dist = dist[:, tf.newaxis]
+    samples = self.evaluate(dist.sample())
+    self.assertAllEqual((6, 4), samples.shape)
+    lp, sliced_lp = self.evaluate([
+        dist.log_prob(samples), sliced_dist.log_prob(samples[:, tf.newaxis])])
+    self.assertAllEqual(lp[:, tf.newaxis], sliced_lp)
+
 
 class EmpiricalScalarStaticShapeTest(
     EmpiricalScalarTest, tf.test.TestCase):