code: def _random_gamma_cpu(
      shape,
      concentration,
      rate=None,
      log_rate=None,
      seed=None,
      log_space=False,
  ):
      """Sample using *fast* `tf.random.stateless_gamma`."""
      bad_concentration = (concentration < 0.0) | tf.math.is_nan(
          concentration
      )
      safe_concentration = tf.where(
          bad_concentration,
          dtype_util.as_numpy_dtype(concentration.dtype)(100.0),
          concentration,
      )
  
      if rate is None:
          if log_rate is None:
              rate = tf.ones([], concentration.dtype)
              log_rate = tf.zeros([], concentration.dtype)
          else:
              rate = tf.math.exp(log_rate)
  
      bad_rate = (rate <= 0.0) | tf.math.is_nan(rate)
  
      if log_space:
          # The underlying gamma sampler uses a recurrence for conc < 1.  When
          # a ~ gamma(conc + 1) and x ~ uniform(0, 1), we have
          #   b = a * x ** (1/conc) ~ gamma(conc)
          # Given that we want log(b) anyway, it's more accurate to just ask the
          # sampler for a (by passing conc + 1 to it in the first place) and
          # do the correction in log-space below.
          orig_safe_concentration = safe_concentration
          safe_concentration = tf.where(
              orig_safe_concentration < 1,
              orig_safe_concentration + 1.0,
              orig_safe_concentration,
          )
          seed, conc_fix_seed = samplers.split_seed(seed)
          log_rate = tf.math.log(rate) if log_rate is None else log_rate
          rate = tf.ones_like(
              log_rate
          )  # Do the division later in log-space.
  
      safe_rate = tf.where(
          bad_rate,
          dtype_util.as_numpy_dtype(concentration.dtype)(100.0),
          rate,
      )
      samples = tf.random.stateless_gamma(
          shape=shape,
          seed=seed,
          alpha=safe_concentration,
          beta=safe_rate,
          dtype=concentration.dtype,
      )
  
      if log_space:
          # Apply the concentration < 1 recurrence here, in log-space.
          samples = tf.math.log(samples)
          conc_fix_unif = samplers.uniform(  # in [0, 1)
              shape, dtype=samples.dtype, seed=conc_fix_seed
          )
  
          conc_lt_one_fix = tf.where(
              orig_safe_concentration < 1,
              # Why do we use log1p(-x)? x is in [0, 1) and log(0) = -inf, is bad.
              # x ~ U(0,1) => 1-x ~ U(0,1)
              # But at the boundary, 1-x in (0, 1]. Good.
              # So we can take log(unif(0,1)) safely as log(1-unif(0,1)).
              # log1p(-0) = 0, and log1p(-almost_one) = -not_quite_inf. Good.
              tf.math.log1p(-conc_fix_unif) / orig_safe_concentration,
              tf.zeros((), dtype=samples.dtype),
          )
          samples += conc_lt_one_fix - log_rate
  
      # 0 rate is infinite scale, which implies a +inf sample.
      # `if log_space` clobbered the `rate` variable with 1 a score lines ago.
      return tf.where(
          (log_rate <= -np.inf if log_space else tf.equal(rate, 0.0)),
          tf.constant(np.inf, dtype=concentration.dtype),
          tf.where(
              bad_rate | bad_concentration,
              dtype_util.as_numpy_dtype(concentration.dtype)(np.nan),
              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.distributions import kullback_leibler
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import batched_rejection_sampler
  from tensorflow_probability.python.internal import custom_gradient
  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 generic
  from tensorflow_probability.python.math import special
  from tensorflow_probability.python.util.deferred_tensor import DeferredTensor
label_desc: Use image node_id 4 for calling a global function with example usage: _random_gamma_cpu(shape, concentration, rate, log_rate, seed, log_space) and returns: tf