code: def _random_binomial_noncpu(
      shape,
      counts,
      probs,
      output_dtype=tf.float32,
      seed=None,
      name=None,
  ):
      """Sample using XLA-friendly python-based rejection sampler."""
      with tf.name_scope(name or "binomial_noncpu"):
          probs = tf.where(counts > 0, probs, 0)
  
          batch_shape = ps.broadcast_shape(
              ps.shape(counts), ps.shape(probs)
          )
          full_shape = ps.concat([shape, batch_shape], axis=0)
  
          inversion_seed, btrs_seed = samplers.split_seed(seed)
  
          p_lt_half = probs < 0.5
          q = tf.where(p_lt_half, probs, 1 - probs)
          q_is_nan = tf.math.is_nan(q)
          q_le_0 = q <= 0.0
          q = tf.where(q_is_nan | q_le_0, tf.constant(0.01, q.dtype), q)
  
          use_inversion = counts * q < 10.0
          counts_for_inversion = tf.where(use_inversion, counts, 0)
          inversion_samples = _binomial_inversion(
              counts_for_inversion, q, full_shape, inversion_seed
          )
  
          counts_for_btrs = tf.where(
              use_inversion, tf.constant(10000.0, counts.dtype), counts
          )
          q_for_btrs = tf.where(
              use_inversion, tf.constant(0.5, q.dtype), q
          )
          btrs_samples = _btrs(
              counts_for_btrs, q_for_btrs, full_shape, btrs_seed
          )
  
          samples = tf.where(
              use_inversion, inversion_samples, btrs_samples
          )
          samples = tf.where(
              q_le_0, tf.zeros([], samples.dtype), samples
          )
          samples = tf.where(
              q_is_nan,
              tf.constant(float("nan"), samples.dtype),
              samples,
          )
          samples = tf.where(p_lt_half, samples, counts - samples)
          return tf.stop_gradient(tf.cast(samples, output_dtype))
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import batched_rejection_sampler
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import implementation_selection
  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 special
label_desc: Use image node_id 6 for calling a global function with example usage: _random_binomial_noncpu(shape, counts, probs, output_dtype, seed, name) and returns: tf