code: def _sample_discrete_gaussian_helper(scale, shape, dtype):
      """Draw samples from discrete Gaussian, assuming scale >= 0."""
      scale = tf.cast(scale, tf.int64)
      sq_scale = _int_square(scale)
  
      # Scale for discrete Laplace. The sampling algorithm should be correct
      # for any discrete Laplace scale, and the original paper uses
      # `dlap_scale = floor(scale) + 1`. Here we use `dlap_scale = scale` (where
      # input `scale` is restricted to integers >= 1) to simplify the fraction
      # below. It turns out that for integer scales >= 1, `dlap_scale = scale` gives
      # a good minimum success rate of ~70%, allowing a small oversampling factor.
      dlap_scale = scale
      oversample_factor = 1.5
  
      # Draw at least some samples in case we got unlucky with small input shape.
      min_n = 1000
      target_n = tf.reduce_prod(tf.cast(shape, tf.int64))
      oversample_n = oversample_factor * tf.cast(target_n, tf.float32)
      draw_n = tf.maximum(min_n, tf.cast(oversample_n, tf.int32))
  
      accepted_n = tf.constant(0, dtype=target_n.dtype)
      result = tf.zeros((0,), dtype=tf.int64)
  
      while accepted_n < target_n:
          # Since the number of samples could be different in every retry, we need to
          # manually specify the shape info for TF.
          tf.autograph.experimental.set_loop_options(
              shape_invariants=[(result, tf.TensorShape([None]))]
          )
          # Draw samples.
          samples = _sample_discrete_laplace(
              dlap_scale, shape=(draw_n,)
          )
          z_numer = _int_square(tf.abs(samples) - scale)
          z_denom = 2 * sq_scale
          bern_probs = tf.exp(-1.0 * tf.divide(z_numer, z_denom))
          accept = _sample_bernoulli(bern_probs)
          # Keep successful samples and increment counter.
          accepted_samples = samples[tf.equal(accept, 1)]
          accepted_n += tf.cast(
              tf.size(accepted_samples), accepted_n.dtype
          )
          result = tf.concat([result, accepted_samples], axis=0)
          # Reduce the number of draws for any retries.
          draw_n = (
              tf.cast(target_n - accepted_n, tf.float32)
              * oversample_factor
          )
          draw_n = tf.maximum(min_n, tf.cast(draw_n, tf.int32))
  
      return tf.cast(tf.reshape(result[:target_n], shape), dtype)
  
imports: import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 5 for calling a global function with example usage: _sample_discrete_gaussian_helper(scale, shape, dtype) and returns: tf