code: def batched_rejection_sampler(
      proposal_fn, target_fn, seed=None, dtype=tf.float32, name=None
  ):
      """Generic batched rejection sampler.
  
      In each iteration, the sampler generates a batch of proposed samples and
      proposal heights by calling `proposal_fn`. For each such sample point S, a
      uniform random variate U is generated and the sample is accepted if U *
      height(S) <= target(S). The batched rejection sampler keeps track of which
      batch points have been accepted, and continues generating new batches until
      all batch points have been acceped.
  
      The values returned by `proposal_fn` should satisfy the desiderata of
      rejection sampling: proposed samples should be drawn independently from a
      valid distribution on some domain D that includes the domain of `target_fn`,
      and the proposal must upper bound the target: for all points S in D, height(S)
      >= target(S).
  
      Args:
        proposal_fn: A callable that takes a Python integer PRNG seed and returns a
          set of proposed samples and the value of the proposal at the samples.
        target_fn: A callable that takes a tensor of samples and returns the value
          of the target at the samples.
        seed: Python integer or `Tensor`, for seeding PRNG.
        dtype: The TensorFlow dtype used internally by `proposal_fn` and
          `target_fn`.  Default value: `tf.float32`.
        name: A name to prepend to created ops.
          Default value: `'batched_rejection_sampler'`.
  
      Returns:
        results, num_iters: A tensor of samples from the target and a scalar int32
        tensor, the number of calls to `proposal_fn`.
      """
      with tf.name_scope(name or "batched_rejection_sampler"):
  
          def randomized_computation(seed):
              """Internal randomized computation."""
              proposal_seed, mask_seed = samplers.split_seed(seed)
  
              proposed_samples, proposed_values = proposal_fn(
                  proposal_seed
              )
  
              # The comparison needs to be strictly less to avoid spurious acceptances
              # when the uniform samples exactly 0 (or when the product underflows to
              # 0).
              target_values = target_fn(proposed_samples)
              good_samples_mask = tf.less(
                  proposed_values
                  * samplers.uniform(
                      ps.shape(proposed_samples),
                      seed=mask_seed,
                      dtype=dtype,
                  ),
                  target_values,
              )
  
              # If either the `proposed_value` or the corresponding `target_value` is
              # `nan`, force that `proposed_sample` to `nan` and accept.  Why?
              #
              # - A `nan` would never be accepted, because tf.less must return False
              #   when either argument is `nan`.
              #
              # - If `nan` happens every time (e.g., due to `nan` in the parameters of
              #   the distribution we are trying to sample from), then we should clearly
              #   return `nan` after going around the rejection loop only once, rather
              #   than looping forever.
              #
              # - If `nan` happens only some of the time, it would silently skew the
              #   distribution on results to always reject, because some of those `nan`
              #   values may have stood for proposals that would have been accepted if
              #   we had computed more accurately.  Instead we forward the `nan`
              #   upstream, so the client can fix their proposal or evaluation
              #   functions.
              #
              # - We force the `proposed_sample` to `nan` because not doing so would
              #   hide a `nan` that occurred in only the `proposed_value` or
              #   `target_value`, silently skewing the distribution on results.
              #
              # - Corner case: if the `proposed_sample` is `nan` but both the
              #   corresponding `proposed_value` and `proposed_target` are for some
              #   reason not `nan`, we trust the user and proceed normally.
              nans = tf.math.is_nan(proposed_values) | tf.math.is_nan(
                  target_values
              )
              proposed_samples = tf.where(
                  nans,
                  tf.cast(np.nan, proposed_samples.dtype),
                  proposed_samples,
              )
              good_samples_mask |= nans
              return proposed_samples, good_samples_mask
  
          samples, _, num_iters = batched_las_vegas_algorithm(
              randomized_computation, seed=seed
          )
          return samples, 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 2 for calling a global function with example usage: batched_rejection_sampler(proposal_fn, target_fn, seed, dtype, name) and returns: samples,  num_iters, proposed_samples,  good_samples_mask