Don't require Root in JDC when sampling with a trivial sample shape.

Pros: In JAX world we often wrap large chunks of computation inside a vmap, which can encompass the log_prob/sample calls of a JDC. In this setting a JDC will never see a non-trivial sample shape. Root was introduced to handle non-trivial sample shapes, so in this setting, this annotation is superfluous. Thus, this change removes the requirement for Root when only trivial sample shapes are considered, improving the UX of JDCs in JAX.

Cons: The check is deferred to non-trivial sample shapes, meaning that a TF user might construct a malformed distribution and not receive an error until they try to sample with a non-trivial sample shape. Prior to this change, they would get an error as soon as `_get_single_sample_distributions` was called, which happens with most property accesses.

Backwards compatibility: Existing well-formed JDCs will continue to work, and will continue being checked for correctness modulo the 'con' above.

PiperOrigin-RevId: 385498069

Author: SiegeLordEx

tensorflow_probability/python/distributions/joint_distribution.py

@@ -787,9 +787,12 @@ def _execute_model(self,
     gen = self._model_coroutine()
     index = 0
     d = next(gen)
-    if self._require_root and not isinstance(d, self.Root):
-      raise ValueError('First distribution yielded by coroutine must '
-                       'be wrapped in `Root`.')
+    if self._require_root:
+      if distribution_util.shape_may_be_nontrivial(
+          sample_shape) and not isinstance(d, self.Root):
+        raise ValueError('First distribution yielded by coroutine must '
+                         'be wrapped in `Root` when requesting a nontrivial '
+                         f'sample_shape = {sample_shape}.')
     try:
       while True:
         actual_distribution = d.distribution if isinstance(d, self.Root) else d

tensorflow_probability/python/distributions/joint_distribution_coroutine.py

@@ -78,7 +78,8 @@ class JointDistributionCoroutine(joint_distribution_lib.JointDistribution):
   The distributions that have been wrapped in the
   `JointDistributionCoroutine.Root` class will be called with `sample_shape` as
   the `sample_shape` argument, and the unwrapped distributions
-  will be called with `()` as the `sample_shape` argument.
+  will be called with `()` as the `sample_shape` argument. The `Root` annotation
+  can be omitted if you never intend to use a `sample_shape` other than `()`.
 
   It is the user's responsibility to ensure that
   each of the distributions generates samples with the specified sample

tensorflow_probability/python/distributions/joint_distribution_coroutine_test.py

@@ -86,7 +86,11 @@ def singleton_jdn_model_fn():
 @test_util.test_all_tf_execution_regimes
 class JointDistributionCoroutineTest(test_util.TestCase):
 
-  def test_batch_and_event_shape_no_plate(self):
+  @parameterized.named_parameters(
+      ('_with_root', True),
+      ('_without_root', False),
+  )
+  def test_batch_and_event_shape_no_plate(self, use_root):
     # The joint distribution specified below corresponds to this
     # graphical model
     #
@@ -95,8 +99,10 @@ def test_batch_and_event_shape_no_plate(self):
     #     \     v
     #      `->-(c)
 
+    root = Root if use_root else lambda x: x
+
     def dist():
-      a = yield Root(tfd.Bernoulli(probs=0.5,
+      a = yield root(tfd.Bernoulli(probs=0.5,
                                    dtype=tf.float32))
       b = yield tfd.Bernoulli(probs=0.25 + 0.5*a,
                               dtype=tf.float32)
@@ -135,7 +141,11 @@ def dist():
     self.assertAllEqual(batch_shape[1], [])
     self.assertAllEqual(batch_shape[2], [])
 
-  def test_batch_and_event_shape_with_plate(self):
+  @parameterized.named_parameters(
+      ('_with_root', True),
+      ('_without_root', False),
+  )
+  def test_batch_and_event_shape_with_plate(self, use_root):
     # The joint distribution specified below corresponds to this
     # graphical model
     #
@@ -146,9 +156,11 @@ def test_batch_and_event_shape_with_plate(self):
     # (df)--+-->---(x)  |
     #       +--------20-+
 
+    root = Root if use_root else lambda x: x
+
     def dist():
-      g = yield Root(tfd.Gamma(2, 2))
-      df = yield Root(tfd.Exponential(1.))
+      g = yield root(tfd.Gamma(2, 2))
+      df = yield root(tfd.Exponential(1.))
       loc = yield tfd.Sample(tfd.Normal(0, g), 20)
       yield tfd.Independent(tfd.StudentT(tf.expand_dims(df, -1), loc, 1), 1)
 
@@ -395,7 +407,7 @@ def dist():
     with self.assertRaisesRegexp(
         Exception,
         'must be wrapped in `Root`'):
-      self.evaluate(joint.sample(seed=test_util.test_seed()))
+      self.evaluate(joint.sample(2, seed=test_util.test_seed()))
 
   @parameterized.named_parameters(
       ('basic', basic_model_with_names_fn),