code: def test_iid_sample_stateful(self):
      # Random fn using stateful samplers.
      def fn(key1, key2, seed=None):
          return [
              normal.Normal(0.0, 1.0).sample([3, 2], seed=seed),
              {
                  key1: poisson.Poisson([1.0, 2.0, 3.0, 4.0]).sample(
                      seed=seed + 1
                  ),
                  key2: lognormal.LogNormal(0.0, 1.0).sample(
                      seed=seed + 2
                  ),
              },
          ]
  
      sample = self.evaluate(
          fn(
              "a",
              key2="b",
              seed=test_util.test_seed(sampler_type="stateful"),
          )
      )
  
      sample_shape = [6, 1]
      iid_fn = vectorization_util.iid_sample(
          fn, sample_shape=sample_shape
      )
      iid_sample = self.evaluate(iid_fn("a", key2="b", seed=42))
  
      # Check that we did not get repeated samples.
      first_sampled_vector = iid_sample[0].flatten()
      self.assertAllGreater(
          (first_sampled_vector[1:] - first_sampled_vector[0]) ** 2,
          1e-6,
      )
  
      expected_iid_shapes = tf.nest.map_structure(
          lambda x: np.concatenate([sample_shape, x.shape], axis=0),
          sample,
      )
      iid_shapes = tf.nest.map_structure(lambda x: x.shape, iid_sample)
      self.assertAllEqualNested(expected_iid_shapes, iid_shapes)
  
imports: import warnings
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import vectorization_util
label_desc: Use image node_id 1 for calling the VectorizationTest obj's underlying member method code with example usage: obj.test_iid_sample_stateful() and returns: list