def _sample_n(self, n, seed=None):
     a = array_ops.ones_like(self.a_b_sum, dtype=self.dtype) * self.a
     b = array_ops.ones_like(self.a_b_sum, dtype=self.dtype) * self.b
     gamma1_sample = random_ops.random_gamma([n], a, dtype=self.dtype, seed=seed)
     gamma2_sample = random_ops.random_gamma([n], b, dtype=self.dtype, seed=seed)
     beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
     return beta_sample

 def loop_fn(i):
   alphas_i = array_ops.gather(alphas, i)
   # Test both scalar and non-scalar params and shapes.
   return (random_ops.random_gamma(alpha=alphas_i[0, 0], shape=[]),
           random_ops.random_gamma(alpha=alphas_i, shape=[]),
           random_ops.random_gamma(alpha=alphas_i[0, 0], shape=[3]),
           random_ops.random_gamma(alpha=alphas_i, shape=[3]))

  def sample_n(self, n, seed=None, name="sample_n"):
    """Sample `n` observations from the Beta Distributions.

    Args:
      n: `Scalar` `Tensor` of type `int32` or `int64`, the number of
        observations to sample.
      seed: Python integer, the random seed.
      name: The name to give this op.

    Returns:
      samples: `[n, ...]`, a `Tensor` of `n` samples for each
        of the distributions determined by broadcasting the hyperparameters.
    """
    with ops.name_scope(self.name):
      with ops.name_scope(name, values=[self.a, self.b, n]):
        a = array_ops.ones_like(self._a_b_sum, dtype=self.dtype) * self.a
        b = array_ops.ones_like(self._a_b_sum, dtype=self.dtype) * self.b
        n = ops.convert_to_tensor(n, name="n")

        gamma1_sample = random_ops.random_gamma(
            [n,], a, dtype=self.dtype, seed=seed)
        gamma2_sample = random_ops.random_gamma(
            [n,], b, dtype=self.dtype, seed=seed)

        beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)

        n_val = tensor_util.constant_value(n)
        final_shape = tensor_shape.vector(n_val).concatenate(
            self._a_b_sum.get_shape())

        beta_sample.set_shape(final_shape)
        return beta_sample

 def _sample_n(self, n, seed=None):
   a = array_ops.ones_like(self.a_b_sum, dtype=self.dtype) * self.a
   b = array_ops.ones_like(self.a_b_sum, dtype=self.dtype) * self.b
   gamma1_sample = random_ops.random_gamma(
       [n,], a, dtype=self.dtype, seed=seed)
   gamma2_sample = random_ops.random_gamma(
       [n,], b, dtype=self.dtype,
       seed=distribution_util.gen_new_seed(seed, "beta"))
   beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
   return beta_sample

 def testShape(self):
   # Fully known shape.
   rnd = random_ops.random_gamma([150], 2.0)
   self.assertEqual([150], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma([150], 2.0, beta=[3.0, 4.0])
   self.assertEqual([150, 2], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma([150], array_ops.ones([1, 2, 3]))
   self.assertEqual([150, 1, 2, 3], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma([20, 30], array_ops.ones([1, 2, 3]))
   self.assertEqual([20, 30, 1, 2, 3], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma(
       [123], array_ops.placeholder(
           dtypes.float32, shape=(2,)))
   self.assertEqual([123, 2], rnd.get_shape().as_list())
   # Partially known shape.
   rnd = random_ops.random_gamma(
       array_ops.placeholder(
           dtypes.int32, shape=(1,)), array_ops.ones([7, 3]))
   self.assertEqual([None, 7, 3], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma(
       array_ops.placeholder(
           dtypes.int32, shape=(3,)), array_ops.ones([9, 6]))
   self.assertEqual([None, None, None, 9, 6], rnd.get_shape().as_list())
   # Unknown shape.
   rnd = random_ops.random_gamma(
       array_ops.placeholder(dtypes.int32),
       array_ops.placeholder(dtypes.float32))
   self.assertIs(None, rnd.get_shape().ndims)
   rnd = random_ops.random_gamma([50], array_ops.placeholder(dtypes.float32))
   self.assertIs(None, rnd.get_shape().ndims)

  def testNoCSE(self):
    """CSE = constant subexpression eliminator.

    SetIsStateful() should prevent two identical random ops from getting
    merged.
    """
    for dtype in dtypes.float16, dtypes.float32, dtypes.float64:
      for use_gpu in [False, True]:
        with self.cached_session(use_gpu=use_gpu):
          rnd1 = random_ops.random_gamma([24], 2.0, dtype=dtype)
          rnd2 = random_ops.random_gamma([24], 2.0, dtype=dtype)
          diff = rnd2 - rnd1
          self.assertGreater(np.linalg.norm(diff.eval()), 0.1)

 def _sample_n(self, n, seed=None):
   """See the documentation for tf.random_gamma for more details."""
   return random_ops.random_gamma([n],
                                  self.alpha,
                                  beta=self.beta,
                                  dtype=self.dtype,
                                  seed=seed)

 def _sample_n(self, n, seed=None):
   n_draws = math_ops.cast(self.n, dtype=dtypes.int32)
   if self.n.get_shape().ndims is not None:
     if self.n.get_shape().ndims != 0:
       raise NotImplementedError(
           "Sample only supported for scalar number of draws.")
   elif self.validate_args:
     is_scalar = check_ops.assert_rank(
         n_draws, 0,
         message="Sample only supported for scalar number of draws.")
     n_draws = control_flow_ops.with_dependencies([is_scalar], n_draws)
   k = self.event_shape()[0]
   unnormalized_logits = array_ops.reshape(
       math_ops.log(random_ops.random_gamma(
           shape=[n],
           alpha=self.alpha,
           dtype=self.dtype,
           seed=seed)),
       shape=[-1, k])
   draws = random_ops.multinomial(
       logits=unnormalized_logits,
       num_samples=n_draws,
       seed=distribution_util.gen_new_seed(seed, salt="dirichlet_multinomial"))
   x = math_ops.reduce_sum(array_ops.one_hot(draws, depth=k),
                           reduction_indices=-2)
   final_shape = array_ops.concat([[n], self.batch_shape(), [k]], 0)
   return array_ops.reshape(x, final_shape)

 def testPositive(self):
   n = int(10e3)
   for dt in [dtypes.float16, dtypes.float32, dtypes.float64]:
     with self.cached_session():
       x = random_ops.random_gamma(shape=[n], alpha=0.001, dtype=dt, seed=0)
       self.assertEqual(0, math_ops.reduce_sum(math_ops.cast(
           math_ops.less_equal(x, 0.), dtype=dtypes.int64)).eval())

  def _sample_n(self, n, seed):
    batch_shape = self.batch_shape_tensor()
    event_shape = self.event_shape_tensor()
    batch_ndims = array_ops.shape(batch_shape)[0]

    ndims = batch_ndims + 3  # sample_ndims=1, event_ndims=2
    shape = array_ops.concat([[n], batch_shape, event_shape], 0)

    # Complexity: O(nbk**2)
    x = random_ops.random_normal(shape=shape,
                                 mean=0.,
                                 stddev=1.,
                                 dtype=self.dtype,
                                 seed=seed)

    # Complexity: O(nbk)
    # This parametrization is equivalent to Chi2, i.e.,
    # ChiSquared(k) == Gamma(alpha=k/2, beta=1/2)
    expanded_df = self.df * array_ops.ones(
        self.scale_operator.batch_shape_tensor(),
        dtype=self.df.dtype.base_dtype)
    g = random_ops.random_gamma(shape=[n],
                                alpha=self._multi_gamma_sequence(
                                    0.5 * expanded_df, self.dimension),
                                beta=0.5,
                                dtype=self.dtype,
                                seed=distribution_util.gen_new_seed(
                                    seed, "wishart"))

    # Complexity: O(nbk**2)
    x = array_ops.matrix_band_part(x, -1, 0)  # Tri-lower.

    # Complexity: O(nbk)
    x = array_ops.matrix_set_diag(x, math_ops.sqrt(g))

    # Make batch-op ready.
    # Complexity: O(nbk**2)
    perm = array_ops.concat([math_ops.range(1, ndims), [0]], 0)
    x = array_ops.transpose(x, perm)
    shape = array_ops.concat([batch_shape, [event_shape[0]], [-1]], 0)
    x = array_ops.reshape(x, shape)

    # Complexity: O(nbM) where M is the complexity of the operator solving a
    # vector system. E.g., for LinearOperatorDiag, each matmul is O(k**2), so
    # this complexity is O(nbk**2). For LinearOperatorLowerTriangular,
    # each matmul is O(k^3) so this step has complexity O(nbk^3).
    x = self.scale_operator.matmul(x)

    # Undo make batch-op ready.
    # Complexity: O(nbk**2)
    shape = array_ops.concat([batch_shape, event_shape, [n]], 0)
    x = array_ops.reshape(x, shape)
    perm = array_ops.concat([[ndims - 1], math_ops.range(0, ndims - 1)], 0)
    x = array_ops.transpose(x, perm)

    if not self.cholesky_input_output_matrices:
      # Complexity: O(nbk^3)
      x = math_ops.matmul(x, x, adjoint_b=True)

    return x

  def _testCompareToExplicitDerivative(self, dtype):
    """Compare to the explicit reparameterization derivative.

    Verifies that the computed derivative satisfies
    dsample / dalpha = d igammainv(alpha, u) / dalpha,
    where u = igamma(alpha, sample).

    Args:
      dtype: TensorFlow dtype to perform the computations in.
    """
    delta = 1e-3
    np_dtype = dtype.as_numpy_dtype
    try:
      from scipy import misc  # pylint: disable=g-import-not-at-top
      from scipy import special  # pylint: disable=g-import-not-at-top

      alpha_val = np.logspace(-2, 3, dtype=np_dtype)
      alpha = constant_op.constant(alpha_val)
      sample = random_ops.random_gamma([], alpha, np_dtype(1.0), dtype=dtype)
      actual = gradients_impl.gradients(sample, alpha)[0]

      (sample_val, actual_val) = self.evaluate((sample, actual))

      u = special.gammainc(alpha_val, sample_val)
      expected_val = misc.derivative(
          lambda alpha_prime: special.gammaincinv(alpha_prime, u),
          alpha_val, dx=delta * alpha_val)

      self.assertAllClose(actual_val, expected_val, rtol=1e-3, atol=1e-3)
    except ImportError as e:
      tf_logging.warn("Cannot use special functions in a test: %s" % str(e))

 def _sample_n(self, n, seed=None):
   gamma_sample = random_ops.random_gamma(
       shape=[n],
       alpha=self.concentration,
       dtype=self.dtype,
       seed=seed)
   return gamma_sample / math_ops.reduce_sum(gamma_sample, -1, keep_dims=True)

 def _sample_n(self, n, seed=None):
   return 1. / random_ops.random_gamma(
       shape=[n],
       alpha=self.concentration,
       beta=self.rate,
       dtype=self.dtype,
       seed=seed)

  def testQuadraticLoss(self):
    """Statistical test for the gradient.

    The equation (5) of https://arxiv.org/abs/1805.08498 says
      d/dalpha E_{sample ~ Gamma(alpha, 1)} f(sample)
        = E_{sample ~ Gamma(alpha, 1)} df(sample)/dalpha.

    Choose a quadratic loss function f(sample) = (sample - t)^2.
    Then, the lhs can be computed analytically:
      d/dalpha E_{sample ~ Gamma(alpha, 1)} f(sample)
        = d/dalpha [ (alpha + alpha^2) - 2 * t * alpha + t^2 ]
        = 1 + 2 * alpha - 2 * t.

    We compare the Monte-Carlo estimate of the expectation with the
    true gradient.
    """
    num_samples = 1000
    t = 0.3
    alpha = 0.5
    expected = 1 + 2 * alpha - 2 * t

    alpha = constant_op.constant(alpha)
    sample = random_ops.random_gamma([num_samples], alpha, 1.0)
    loss = math_ops.reduce_mean(math_ops.square(sample - t))
    dloss_dalpha = gradients_impl.gradients(loss, alpha)[0]
    dloss_dalpha_val = self.evaluate(dloss_dalpha)
    self.assertAllClose(expected, dloss_dalpha_val, atol=1e-1, rtol=1e-1)

 def _sample_n(self, n, seed=None):
   expanded_concentration1 = array_ops.ones_like(
       self.total_concentration, dtype=self.dtype) * self.concentration1
   expanded_concentration0 = array_ops.ones_like(
       self.total_concentration, dtype=self.dtype) * self.concentration0
   gamma1_sample = random_ops.random_gamma(
       shape=[n],
       alpha=expanded_concentration1,
       dtype=self.dtype,
       seed=seed)
   gamma2_sample = random_ops.random_gamma(
       shape=[n],
       alpha=expanded_concentration0,
       dtype=self.dtype,
       seed=distribution_util.gen_new_seed(seed, "beta"))
   beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
   return beta_sample

 def func():
   with self.session(use_gpu=use_gpu, graph=ops.Graph()) as sess:
     rng = random_ops.random_gamma(
         [num], alpha, beta=beta, dtype=dtype, seed=seed)
     ret = np.empty([10, num])
     for i in xrange(10):
       ret[i, :] = sess.run(rng)
   return ret

 def testGradientsShape(self):
   shape = [2, 3]
   alpha = array_ops.ones([2, 2])
   beta = array_ops.ones([1, 2])
   sample = random_ops.random_gamma(shape, alpha, beta, seed=12345)
   grads_alpha, grads_beta = gradients_impl.gradients(sample, [alpha, beta])
   self.assertAllEqual(grads_alpha.shape, alpha.shape)
   self.assertAllEqual(grads_beta.shape, beta.shape)

 def testGradientsShapeWithOneSamplePerParameter(self):
   shape = []
   alpha = array_ops.ones([2, 2])
   beta = array_ops.ones([1, 2])
   sample = random_ops.random_gamma(shape, alpha, beta)
   grads_alpha, grads_beta = gradients_impl.gradients(sample, [alpha, beta])
   self.assertAllEqual(grads_alpha.shape, alpha.shape)
   self.assertAllEqual(grads_beta.shape, beta.shape)

  def _sample_n(self, n, seed):
    batch_shape = self.batch_shape()
    event_shape = self.event_shape()
    batch_ndims = array_ops.shape(batch_shape)[0]

    ndims = batch_ndims + 3  # sample_ndims=1, event_ndims=2
    shape = array_ops.concat(((n,), batch_shape, event_shape), 0)

    # Complexity: O(nbk^2)
    x = random_ops.random_normal(shape=shape,
                                 mean=0.,
                                 stddev=1.,
                                 dtype=self.dtype,
                                 seed=seed)

    # Complexity: O(nbk)
    # This parametrization is equivalent to Chi2, i.e.,
    # ChiSquared(k) == Gamma(alpha=k/2, beta=1/2)
    g = random_ops.random_gamma(shape=(n,),
                                alpha=self._multi_gamma_sequence(
                                    0.5 * self.df, self.dimension),
                                beta=0.5,
                                dtype=self.dtype,
                                seed=distribution_util.gen_new_seed(
                                    seed, "wishart"))

    # Complexity: O(nbk^2)
    x = array_ops.matrix_band_part(x, -1, 0)  # Tri-lower.

    # Complexity: O(nbk)
    x = array_ops.matrix_set_diag(x, math_ops.sqrt(g))

    # Make batch-op ready.
    # Complexity: O(nbk^2)
    perm = array_ops.concat((math_ops.range(1, ndims), (0,)), 0)
    x = array_ops.transpose(x, perm)
    shape = array_ops.concat((batch_shape, (event_shape[0], -1)), 0)
    x = array_ops.reshape(x, shape)

    # Complexity: O(nbM) where M is the complexity of the operator solving a
    # vector system.  E.g., for OperatorPDDiag, each matmul is O(k^2), so
    # this complexity is O(nbk^2). For OperatorPDCholesky, each matmul is
    # O(k^3) so this step has complexity O(nbk^3).
    x = self.scale_operator_pd.sqrt_matmul(x)

    # Undo make batch-op ready.
    # Complexity: O(nbk^2)
    shape = array_ops.concat((batch_shape, event_shape, (n,)), 0)
    x = array_ops.reshape(x, shape)
    perm = array_ops.concat(((ndims - 1,), math_ops.range(0, ndims - 1)), 0)
    x = array_ops.transpose(x, perm)

    if not self.cholesky_input_output_matrices:
      # Complexity: O(nbk^3)
      x = math_ops.matmul(x, x, adjoint_b=True)

    return x

 def _sample_n(self, n, seed=None):
   # The sampling method comes from the well known fact that if X ~ Normal(0,
   # 1), and Z ~ Chi2(df), then X / sqrt(Z / df) ~ StudentT(df).
   shape = array_ops.concat(0, ([n], self.batch_shape()))
   normal_sample = random_ops.random_normal(
       shape, dtype=self.dtype, seed=seed)
   half = constant_op.constant(0.5, self.dtype)
   df = self.df * array_ops.ones(self.batch_shape(), dtype=self.dtype)
   gamma_sample = random_ops.random_gamma(
       [n,], half * df, beta=half, dtype=self.dtype,
       seed=distribution_util.gen_new_seed(seed, salt="student_t"))
   samples = normal_sample / math_ops.sqrt(gamma_sample / df)
   return samples * self.sigma + self.mu