code: def _sample_multinomial_as_iterated_binomial(
      num_samples, num_classes, probs, num_trials, dtype, seed
  ):
      """Sample a multinomial by drawing one binomial sample per class.
  
      The batch shape is given by broadcasting num_trials with
      remove_last_dimension(probs).
  
      The loop over binomial samples is a `tf.while_loop`, thus supporting a dynamic
      number of classes.
  
      Args:
        num_samples: Singleton integer Tensor: number of multinomial samples to
          draw.
        num_classes: Singleton integer Tensor: number of classes.
        probs: Floating Tensor with last dimension `num_classes`, of normalized
          probabilities per class.
        num_trials: Tensor of number of categorical trials each multinomial consists
          of.  num_trials[..., tf.newaxis] must broadcast with probs.
        dtype: dtype at which to emit samples.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
      Returns:
        samples: Tensor of given dtype and shape [num_samples] + batch_shape +
          [num_classes].
      """
      with tf.name_scope("draw_sample"):
          # `convert_to_tensor(num_classes) here to avoid unstacking inside
          # `split_seed`.  We can't take advantage of the Python-list code path anyway
          # because the index at which we will take the seed is a Tensor.
          seeds = samplers.split_seed(
              seed,
              n=ps.convert_to_shape_tensor(num_classes),
              salt="multinomial_draw_sample",
          )
  
          def fn(i, num_trials, consumed_prob, accum):
              """Sample the counts for one class using binomial."""
              probs_here = tf.gather(probs, i, axis=-1)
              binomial_probs = tf.clip_by_value(
                  # Divide without nan; we could have probs_here=0 and consumed_prob=1.
                  tf.math.divide_no_nan(
                      probs_here, (1.0 - consumed_prob)
                  ),
                  0,
                  1,
              )
              seed_here = tf.gather(seeds, i, axis=0)
              binom = binomial.Binomial(
                  total_count=num_trials, probs=binomial_probs
              )
              # Not passing `num_samples` to `binom.sample`, as it's is already in
              # `num_trials.shape`.
              sample = binom.sample(seed=seed_here)
              accum = accum.write(i, tf.cast(sample, dtype=dtype))
              return (
                  i + 1,
                  num_trials - sample,
                  consumed_prob + probs_here,
                  accum,
              )
  
          num_trials = tf.cast(num_trials, probs.dtype)
          # Pre-broadcast with probs
          num_trials = num_trials + tf.zeros_like(probs[..., 0])
          # Pre-enlarge for different output samples
          num_trials = _replicate_along_left(num_trials, num_samples)
          i = tf.constant(0)
          consumed_prob = tf.zeros_like(probs[..., 0])
          accum = tf.TensorArray(
              dtype, size=num_classes, element_shape=num_trials.shape
          )
          _, num_trials_left, _, accum = tf.while_loop(
              cond=lambda index, _0, _1, _2: tf.less(
                  index, num_classes - 1
              ),
              body=fn,
              loop_vars=(i, num_trials, consumed_prob, accum),
          )
          # Force the last iteration to put all the trials into the last bucket,
          # because probs[..., -1] / (1. - consumed_prob) might numerically not be 1.
          # Also saves one iteration around the while_loop and one run of the binomial
          # sampler.
          accum = accum.write(
              num_classes - 1, tf.cast(num_trials_left, dtype=dtype)
          )
          # This stop_gradient is necessary to prevent spurious zero gradients coming
          # from b/138796859, and a spurious gradient through num_trials_left.
          results = tf.stop_gradient(accum.stack())
          return distribution_util.move_dimension(results, 0, -1)
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softmax_centered
  from tensorflow_probability.python.distributions import binomial
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 2 for calling a global function with example usage: _sample_multinomial_as_iterated_binomial(num_samples, num_classes, probs, num_trials, dtype, seed) and returns: distribution_util, unknown,  unknown,  unknown,  accum