code: def _random_gamma_rejection(
      shape,
      concentration,
      rate=None,
      log_rate=None,
      seed=None,
      log_space=False,
      internal_dtype=None,
  ):
      """Samples from the gamma distribution.
  
      The sampling algorithm is rejection sampling [1], and pathwise gradients with
      respect to concentration are computed via implicit differentiation [2].
  
      Args:
        shape: The output sample shape. Trailing dims must match broadcast of
          `concentration` with `rate` or `log_rate`.
        concentration: Floating point tensor, the concentration params of the
          distribution(s). Must contain only positive values. Must broadcast with
          `rate` or `log_rate`, if given.
        rate: Floating point tensor, the inverse scale params of the
          distribution(s). Must contain only positive values. Must broadcast with
          `concentration`. If `None`, handled as if 1 (but possibly more
          efficiently). Mutually exclusive with `log_rate`.
        log_rate: Floating point tensor, log of the inverse scale params of the
          distribution(s). Must broadcast with `concentration`. If `None`, handled
          as if 0 (but possibly more efficiently). Mutually exclusive with `rate`.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        log_space: Optionally sample log(gamma) variates.
        internal_dtype: dtype to use for internal computations. If unspecified, we
          use the same dtype as the output (i.e. the dtype of `concentration`,
          `rate`, or `log_rate`) when `log_space==True`, and `tf.float64` otherwise.
  
      Returns:
        Differentiable samples from the gamma distribution.
  
      #### References
  
      [1] George Marsaglia and Wai Wan Tsang. A simple method for generating Gamma
          variables. ACM Transactions on Mathematical Software, 2000.
  
      [2] Michael Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit
          Reparameterization Gradients. Neural Information Processing Systems, 2018.
      """
      (
          generate_and_test_samples_seed,
          concentration_fix_seed,
      ) = samplers.split_seed(seed, salt="random_gamma")
      output_dtype = dtype_util.common_dtype(
          [concentration, rate, log_rate], dtype_hint=tf.float32
      )
      if internal_dtype is None:
          internal_dtype = output_dtype if log_space else tf.float64
  
      def rejection_sample(concentration):
          """Gamma rejection sampler."""
          # Note, concentration here already has a shape that is broadcast with rate.
          cast_concentration = tf.cast(concentration, internal_dtype)
  
          good_params_mask = concentration >= 0.0
          # When replacing NaN values, use 100. for concentration, since that leads to
          # a high-likelihood of the rejection sampler accepting on the first pass.
          safe_concentration = tf.where(
              good_params_mask, cast_concentration, 100.0
          )
  
          modified_safe_concentration = tf.where(
              safe_concentration < 1.0,
              safe_concentration + 1.0,
              safe_concentration,
          )
  
          one_third = tf.constant(1.0 / 3, dtype=internal_dtype)
          d = modified_safe_concentration - one_third
          c = one_third * tf.math.rsqrt(d)
  
          def generate_and_test_samples(seed):
              """Generate and test samples."""
              v_seed, u_seed = samplers.split_seed(seed)
  
              x = samplers.normal(
                  shape, dtype=internal_dtype, seed=v_seed
              )
              # This implicitly broadcasts concentration up to sample shape.
              v = 1 + c * x
              # In [1], there is an 'inner' rejection sampling loop which checks that
              # v > 0 and generates a new normal sample if it's not, saving the rest of
              # the computations below. We found that merging the check for  v > 0 with
              # the `good_sample_mask` not only simplifies the code, but leads to a
              # ~2x speedup for small concentrations on GPU, at the cost of deviating
              # slightly from the implementation given in Ref. [1].
              accept_v = v > 0.0
              logv = tf.math.log1p(c * x)
              x2 = x * x
              v3 = v * v * v
              logv3 = logv * 3
  
              u = samplers.uniform(
                  shape, dtype=internal_dtype, seed=u_seed
              )
  
              # In [1], the suggestion is to first check u < 1 - 0.331 * x2 * x2, and to
              # run the check below only if it fails, in order to avoid the relatively
              # expensive logarithm calls. Our algorithm operates in batch mode: we will
              # have to compute or not compute the logarithms for the entire batch, and
              # as the batch gets larger, the odds we compute it grow. Therefore we
              # don't bother with the "cheap" check.
              good_sample_mask = tf.logical_and(
                  tf.math.log(u) < (x2 / 2.0 + d * (1 - v3 + logv3)),
                  accept_v,
              )
  
              return logv3 if log_space else v3, good_sample_mask
  
          samples = brs.batched_las_vegas_algorithm(
              generate_and_test_samples,
              seed=generate_and_test_samples_seed,
          )[0]
  
          concentration_fix_unif = samplers.uniform(  # in [0, 1)
              shape, dtype=internal_dtype, seed=concentration_fix_seed
          )
  
          if log_space:
              concentration_lt_one_fix = tf.where(
                  safe_concentration < 1.0,
                  # 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(-concentration_fix_unif)
                  / safe_concentration,
                  tf.zeros((), dtype=internal_dtype),
              )
              samples = (
                  samples + tf.math.log(d) + concentration_lt_one_fix
              )
          else:
              concentration_lt_one_fix = tf.where(
                  safe_concentration < 1.0,
                  tf.math.pow(
                      concentration_fix_unif,
                      tf.math.reciprocal(safe_concentration),
                  ),
                  tf.ones((), dtype=internal_dtype),
              )
              samples = samples * d * concentration_lt_one_fix
  
          samples = tf.where(good_params_mask, samples, np.nan)
          output_type_samples = tf.cast(samples, output_dtype)
  
          return output_type_samples
  
      broadcast_conc_shape = ps.broadcast_shape(
          ps.shape(concentration), _shape_or_scalar(rate, log_rate)
      )
      broadcast_concentration = tf.broadcast_to(
          concentration, broadcast_conc_shape
      )
      concentration_samples = rejection_sample(broadcast_concentration)
  
      if rate is not None and log_rate is not None:
          raise ValueError(
              "`rate` and `log_rate` are mutually exclusive."
          )
  
      if rate is None and log_rate is None:
          if not log_space:
              concentration_samples = _fix_zero_samples(
                  concentration_samples
              )
          return concentration_samples
  
      if log_space:
          if log_rate is None:
              log_rate = tf.math.log(
                  tf.where(rate >= 0.0, rate, np.nan)
              )
          return concentration_samples - log_rate
      else:
          if rate is None:
              rate = tf.math.exp(log_rate)
          corrected_rate = tf.where(rate >= 0.0, rate, np.nan)
          # 0 rate is infinite scale, which implies a +inf sample.
          ret = tf.where(
              tf.equal(corrected_rate, 0),
              tf.constant(np.inf, dtype=output_dtype),
              _fix_zero_samples(concentration_samples / corrected_rate),
          )
          return ret
  
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 15 for calling a global function with example usage: _random_gamma_rejection(shape, concentration, rate, log_rate, seed, log_space, internal_dtype) and returns: output_type_samples, concentration_samples, unknown, ret, unknown,  good_sample_mask