code: def _random_poisson_noncpu(
      shape,
      rates=None,
      log_rates=None,
      output_dtype=tf.float32,
      seed=None,
      name=None,
  ):
      """Sample using XLA-friendly python-based rejection sampler."""
      with tf.name_scope(name or "poisson_noncpu"):
          if log_rates is None:
              log_rates = tf.math.log(rates)
          shape = ps.concat([shape, ps.shape(log_rates)], axis=0)
          good_params_mask = ~tf.math.is_nan(log_rates)
          internal_dtype = tf.float64
  
          seed_lo, seed_hi = samplers.split_seed(seed)
  
          # First, we sample the values for which rate >= 10.
          # When replacing NaN or < 10 values, use 100 for log rate, since that leads
          # to a high-likelihood of the rejection sampler accepting on the first pass.
          high_params_mask = good_params_mask & (
              log_rates >= np.log(10.0)
          )
          cast_log_rates = tf.cast(log_rates, internal_dtype)
          safe_log_rates = tf.where(
              high_params_mask, cast_log_rates, 100.0
          )
          high_rate_samples = _random_poisson_high_rate(
              shape,
              log_rate=safe_log_rates,
              internal_dtype=internal_dtype,
              seed=seed_hi,
          )
          high_rate_samples = tf.cast(high_rate_samples, output_dtype)
  
          # Next, we sample the values for which rate < 10. When replacing NaN or high
          # values, use a small number so that the sum-of-exponentials sampler
          # terminates on the first pass with high likelihood.
          low_params_mask = good_params_mask & ~high_params_mask
          safe_rate = tf.where(
              low_params_mask, tf.math.exp(cast_log_rates), 1e-5
          )
          low_rate_samples = _random_poisson_low_rate(
              shape,
              rate=safe_rate,
              internal_dtype=internal_dtype,
              seed=seed_lo,
          )
          low_rate_samples = tf.cast(low_rate_samples, output_dtype)
  
          samples = tf.where(
              good_params_mask,
              tf.where(
                  high_params_mask, high_rate_samples, low_rate_samples
              ),
              np.nan,
          )
  
      return samples
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  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
label_desc: Use image node_id 2 for calling a global function with example usage: _random_poisson_noncpu(shape, rates, log_rates, output_dtype, seed, name) and returns: samples