Avoid sampling from JDSequential/JDNamed models at init time.

Now that we use static tracing for most JD properties, it makes sense to defer the cost of building and caching a list of distributions until some method actually needs them.

PiperOrigin-RevId: 385001342

Author: davmre

tensorflow_probability/python/distributions/joint_distribution_sequential.py

@@ -210,9 +210,12 @@ def __init__(self, model, validate_args=False, name=None):
       self._build(model)
 
       self._single_sample_distributions = {}
-      self._get_single_sample_distributions(candidate_dists=[
-          None if a else d() for d, a
-          in zip(self._dist_fn_wrapped, self._dist_fn_args)])
+
+      # If the model consists entirely of prebuilt distributions with no
+      # dependencies, cache them directly to avoid a sample call down the road.
+      if not any(self._dist_fn_args):
+        self._get_single_sample_distributions(
+            candidate_dists=[d() for d in self._dist_fn_wrapped])
 
       super(JointDistributionSequential, self).__init__(
           dtype=None,  # Ignored; we'll override.

tensorflow_probability/python/distributions/joint_distribution_sequential_test.py

@@ -176,7 +176,8 @@ def test_norequired_args_maker(self):
     """Test that only non-default args are passed through."""
     with self.assertRaisesWithPredicateMatch(
         ValueError, 'Must pass probs or logits, but not both.'):
-      tfd.JointDistributionSequential([tfd.Normal(0., 1.), tfd.Bernoulli])
+      d = tfd.JointDistributionSequential([tfd.Normal(0., 1.), tfd.Bernoulli])
+      d.sample(seed=test_util.test_seed())
 
   def test_graph_resolution(self):
     d = tfd.JointDistributionSequential(
@@ -767,6 +768,28 @@ def test_creates_valid_coroutine(self):
             jdc.sample([5], seed=test_util.test_seed()))]
     self.assertAllEqualNested(sample_shapes, jdc_sample_shapes)
 
+  def test_init_does_not_execute_model(self):
+    model_traces = []
+    def record_model_called(x):
+      model_traces.append(x)
+      return x
+
+    model = tfd.JointDistributionSequential(
+        [
+            tfd.Normal(0., 1.),
+            lambda z: tfd.Normal(record_model_called(z), 1.)
+        ],
+        validate_args=True)
+    # Model should not be called from init.
+    self.assertLen(model_traces, 0)
+    model.sample(seed=test_util.test_seed())
+    # The first sample call will run the model twice (for shape
+    # inference + actually sampling).
+    self.assertLen(model_traces, 2)
+    # Subsequent calls should only run the model once.
+    model.sample([2], seed=test_util.test_seed())
+    self.assertLen(model_traces, 3)
+
 
 class ResolveDistributionNamesTest(test_util.TestCase):