code: def batched_las_vegas_algorithm(
      batched_las_vegas_trial_fn, max_trials=None, seed=None, name=None
  ):
      """Batched Las Vegas Algorithm.
  
      This utility encapsulates the notion of a 'batched las_vegas_algorithm'
      (BLVA): a batch of independent (but not necessarily identical) randomized
      computations, each of which will eventually terminate after an unknown number
      of trials [(Babai, 1979)][1]. The number of trials will in general vary
      across batch points.
  
      The computation is parameterized by a callable representing a single trial for
      the entire batch. The utility runs the callable repeatedly, keeping track of
      which batch points have succeeded, until all have succeeded.
  
      Because we keep running the callable repeatedly until we've generated at least
      one good value for every batch point, we may generate multiple good values for
      many batch points. In this case, the particular good batch point returned is
      deliberately left unspecified.
  
      Args:
        batched_las_vegas_trial_fn: A callable that takes a PRNG seed and returns
          two values. (1) A structure of Tensors containing the results of the
          computation, all with a shape broadcastable with (2) a boolean mask
          representing whether each batch point succeeded.
        max_trials: An optional integer Tensor giving an upper bound on
          the number of calls to `batched_las_vegas_trial_fn`.  If not
          supplied, no limit is enforced.
        seed: Python integer or `Tensor`, for seeding PRNG.
        name: A name to prepend to created ops.
          Default value: `'batched_las_vegas_algorithm'`.
  
      Returns:
        results: A structure of Tensors representing the results of a
          successful computation for each batch point.
        successes: A boolean Tensor indicating which batch points succeeded.
          This will be all `True` if no `max_trials` limit was supplied.
        num_iters: A scalar int32 tensor, the number of calls to
          `batched_las_vegas_algorithm`.
  
      #### References
  
      [1]: Laszlo Babai. Monte-Carlo algorithms in graph isomorphism
           testing. Universite de Montreal, D.M.S. No. 79-10.
      """
      with tf.name_scope(name or "batched_las_vegas_algorithm"):
          init_seed, loop_seed = samplers.split_seed(seed)
          values, good_values_mask = batched_las_vegas_trial_fn(
              init_seed
          )
          num_iters = tf.constant(1)
  
          def cond(
              unused_values, good_values_mask, num_iters, unused_seed
          ):
              all_done = tf.reduce_all(good_values_mask)
              if max_trials is not None:
                  return ~all_done & (num_iters < max_trials)
              else:
                  return ~all_done
  
          def body(values, good_values_mask, num_iters, seed):
              """Batched Las Vegas Algorithm body."""
  
              trial_seed, new_seed = samplers.split_seed(seed)
              (
                  new_values,
                  new_good_values_mask,
              ) = batched_las_vegas_trial_fn(trial_seed)
              # TODO(b/117888585): Remove this workaround for shape inference issues.
              tensorshape_util.set_shape(
                  new_good_values_mask, good_values_mask.shape
              )
  
              def pick(new, old):
                  ret = bu.where_left_justified_mask(
                      new_good_values_mask, new, old
                  )
                  # TODO(b/117888585): Remove this workaround for shape inference issues.
                  tensorshape_util.set_shape(ret, old.shape)
                  return ret
  
              values = tf.nest.map_structure(pick, new_values, values)
  
              good_values_mask = good_values_mask | new_good_values_mask
  
              return values, good_values_mask, num_iters + 1, new_seed
  
          (values, final_successes, num_iters, _) = tf.while_loop(
              cond,
              body,
              (values, good_values_mask, num_iters, loop_seed),
          )
          values = tf.nest.map_structure(tf.stop_gradient, values)
          return values, final_successes, num_iters
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: batched_las_vegas_algorithm(batched_las_vegas_trial_fn, max_trials, seed, name) and returns: values,  final_successes,  num_iters, values,  good_values_mask,  unknown,  new_seed, unknown, unknown, ret