code: def iid_sample(sample_fn, sample_shape):
      """Lift a sampling function to one that draws multiple iid samples.
  
      Args:
        sample_fn: Python `callable` that returns a (possibly nested) structure of
          `Tensor`s. May optionally take a `seed` named arg: if so, any `int`
          seeds (for stateful samplers) are passed through directly, while any
          pair-of-`int` seeds (for stateless samplers) are split into independent
          seeds for each sample.
        sample_shape: `int` `Tensor` shape of iid samples to draw.
      Returns:
        iid_sample_fn: Python `callable` taking the same arguments as `sample_fn`
          and returning iid samples. Each returned `Tensor` will have shape
          `concat([sample_shape, shape_of_original_returned_tensor])`.
      """
      sample_shape = distribution_util.expand_to_vector(
          ps.cast(sample_shape, np.int32), tensor_name="sample_shape"
      )
      n = ps.cast(ps.reduce_prod(sample_shape), dtype=np.int32)
      static_n = tf.get_static_value(tf.convert_to_tensor(n))
  
      def unflatten(x):
          sample_dims = 0 if static_n == 1 else 1
          unflattened_shape = ps.cast(
              ps.concat(
                  [sample_shape, ps.shape(x)[sample_dims:]], axis=0
              ),
              dtype=np.int32,
          )
          return tf.reshape(x, unflattened_shape)
  
      def iid_sample_fn(*args, **kwargs):
          """Draws iid samples from `fn`."""
  
          with tf.name_scope("iid_sample_fn"):
              seed = kwargs.pop("seed", None)
              if samplers.is_stateful_seed(seed):
                  kwargs = dict(
                      kwargs, seed=SeedStream(seed, salt="iid_sample")()
                  )
  
                  def pfor_loop_body(_):
                      with tf.name_scope("iid_sample_fn_stateful_body"):
                          return sample_fn(*args, **kwargs)
  
              else:
                  # If a stateless seed arg is passed, split it into `n` different
                  # stateless seeds, so that we don't just get a bunch of copies of the
                  # same sample.
                  if not JAX_MODE:
                      warnings.warn(
                          "Saw Tensor seed {}, implying stateless sampling. Autovectorized "
                          "functions that use stateless sampling may be quite slow because "
                          "the current implementation falls back to an explicit loop. This "
                          "will be fixed in the future. For now, you will likely see "
                          "better performance from stateful sampling, which you can invoke "
                          "by passing a Python `int` seed.".format(seed)
                      )
                  seed = samplers.split_seed(
                      seed, n=n, salt="iid_sample_stateless"
                  )
  
                  def pfor_loop_body(i):
                      with tf.name_scope(
                          "iid_sample_fn_stateless_body"
                      ):
                          return sample_fn(
                              *args, seed=tf.gather(seed, i), **kwargs
                          )
  
              if static_n == 1:
                  draws = pfor_loop_body(0)
              else:
                  draws = control_flow_ops.pfor(pfor_loop_body, n)
              return tf.nest.map_structure(
                  unflatten, draws, expand_composites=True
              )
  
      return iid_sample_fn
  
file_path: probability/tensorflow_probability/python/internal/vectorization_util.py
imports: import functools
  import warnings
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.util import SeedStream
  from tensorflow.python.ops.parallel_for import control_flow_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 1 for calling a global function with example usage: iid_sample(sample_fn, sample_shape) and returns: iid_sample_fn, tf, tf, sample_fn, sample_fn