Fix Beta, Wishart and Mixture sampling to use different seeds per random draw.

Change: 136075925

tensorflow/contrib/distributions/python/kernel_tests/beta_test.py

@@ -245,6 +245,23 @@ class BetaTest(tf.test.TestCase):
                           stats.beta.var(a, b),
                           atol=1e-1)
 
+  # Test that sampling with the same seed twice gives the same results.
+  def testBetaSampleMultipleTimes(self):
+    with self.test_session():
+      a_val = 1.
+      b_val = 2.
+      n_val = 100
+
+      tf.set_random_seed(654321)
+      beta1 = tf.contrib.distributions.Beta(a=a_val, b=b_val, name="beta1")
+      samples1 = beta1.sample_n(n_val, seed=123456).eval()
+
+      tf.set_random_seed(654321)
+      beta2 = tf.contrib.distributions.Beta(a=a_val, b=b_val, name="beta2")
+      samples2 = beta2.sample_n(n_val, seed=123456).eval()
+
+      self.assertAllClose(samples1, samples2)
+
   def testBetaSampleMultidimensional(self):
     with self.test_session():
       a = np.random.rand(3, 2, 2).astype(np.float32)

tensorflow/contrib/distributions/python/kernel_tests/mixture_test.py

@@ -334,6 +334,32 @@ class MixtureTest(tf.test.TestCase):
         which_dist_samples = dist_sample_values[c][:size_c]
         self.assertAllClose(which_dist_samples, sample_values[which_c])
 
+  # Test that sampling with the same seed twice gives the same results.
+  def testSampleMultipleTimes(self):
+    # 5 component mixture.
+    logits = [-10.0, -5.0, 0.0, 5.0, 10.0]
+    mus = [-5.0, 0.0, 5.0, 4.0, 20.0]
+    sigmas = [0.1, 5.0, 3.0, 0.2, 4.0]
+
+    with self.test_session():
+      n = 100
+
+      tf.set_random_seed(654321)
+      components = [distributions_py.Normal(
+          mu=mu, sigma=sigma) for mu, sigma in zip(mus, sigmas)]
+      cat = distributions_py.Categorical(logits, dtype=tf.int32, name="cat1")
+      dist1 = distributions_py.Mixture(cat, components, name="mixture1")
+      samples1 = dist1.sample_n(n, seed=123456).eval()
+
+      tf.set_random_seed(654321)
+      components2 = [distributions_py.Normal(
+          mu=mu, sigma=sigma) for mu, sigma in zip(mus, sigmas)]
+      cat2 = distributions_py.Categorical(logits, dtype=tf.int32, name="cat2")
+      dist2 = distributions_py.Mixture(cat2, components2, name="mixture2")
+      samples2 = dist2.sample_n(n, seed=123456).eval()
+
+      self.assertAllClose(samples1, samples2)
+
   def testSampleScalarBatchMultivariate(self):
     with self.test_session() as sess:
       num_components = 3

tensorflow/contrib/distributions/python/kernel_tests/student_t_test.py

@@ -130,16 +130,17 @@ class StudentTTest(tf.test.TestCase):
       mu_v = 3.0
       sigma_v = np.sqrt(10.0)
       n = tf.constant(100)
+
+      tf.set_random_seed(654321)
       student = tf.contrib.distributions.StudentT(
-          name="student_t1", df=df, mu=mu, sigma=sigma)
+          df=df, mu=mu, sigma=sigma, name="student_t1")
       samples1 = student.sample_n(n, seed=123456).eval()
 
-    # We need to start another test session, since this resets the internal
-    # state.
-    with self.test_session():
+      tf.set_random_seed(654321)
       student2 = tf.contrib.distributions.StudentT(
-          name="student_t2", df=df, mu=mu, sigma=sigma)
+          df=df, mu=mu, sigma=sigma, name="student_t2")
       samples2 = student2.sample_n(n, seed=123456).eval()
+
       self.assertAllClose(samples1, samples2)
 
   def testStudentSampleSmallDfNoNan(self):

tensorflow/contrib/distributions/python/kernel_tests/wishart_test.py

@@ -149,6 +149,30 @@ class WishartCholeskyTest(tf.test.TestCase):
                           variance_estimate,
                           rtol=0.05)
 
+  # Test that sampling with the same seed twice gives the same results.
+  def testSampleMultipleTimes(self):
+    with self.test_session():
+      df = 4.
+      n_val = 100
+
+      tf.set_random_seed(654321)
+      chol_w1 = distributions.WishartCholesky(
+          df=df,
+          scale=chol(make_pd(1., 3)),
+          cholesky_input_output_matrices=False,
+          name="wishart1")
+      samples1 = chol_w1.sample_n(n_val, seed=123456).eval()
+
+      tf.set_random_seed(654321)
+      chol_w2 = distributions.WishartCholesky(
+          df=df,
+          scale=chol(make_pd(1., 3)),
+          cholesky_input_output_matrices=False,
+          name="wishart2")
+      samples2 = chol_w2.sample_n(n_val, seed=123456).eval()
+
+      self.assertAllClose(samples1, samples2)
+
   def testProb(self):
     with self.test_session():
       # Generate some positive definite (pd) matrices and their Cholesky

tensorflow/contrib/distributions/python/ops/beta.py

@@ -197,7 +197,8 @@ class Beta(distribution.Distribution):
     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)
+        [n,], b, dtype=self.dtype,
+        seed=distribution_util.gen_new_seed(seed, "beta"))
     beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
     return beta_sample
 

tensorflow/contrib/distributions/python/ops/mixture.py

@@ -22,6 +22,7 @@ import numpy as np
 
 from tensorflow.contrib.distributions.python.ops import categorical
 from tensorflow.contrib.distributions.python.ops import distribution
+from tensorflow.contrib.distributions.python.ops import distribution_util
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_util
@@ -295,8 +296,10 @@ class Mixture(distribution.Distribution):
           partitions=cat_samples,
           num_partitions=self.num_components)
       samples_class = [None for _ in range(self.num_components)]
+
       for c in range(self.num_components):
         n_class = array_ops.size(partitioned_samples_indices[c])
+        seed = distribution_util.gen_new_seed(seed, "mixture")
         samples_class_c = self.components[c].sample_n(n_class, seed=seed)
 
         # Pull out the correct batch entries from each index.

tensorflow/contrib/distributions/python/ops/wishart.py

@@ -22,6 +22,7 @@ import math
 import numpy as np
 
 from tensorflow.contrib.distributions.python.ops import distribution
+from tensorflow.contrib.distributions.python.ops import distribution_util
 from tensorflow.contrib.distributions.python.ops import operator_pd_cholesky
 from tensorflow.contrib.distributions.python.ops import operator_pd_full
 from tensorflow.contrib.framework.python.framework import tensor_util as contrib_tensor_util
@@ -211,7 +212,8 @@ class _WishartOperatorPD(distribution.Distribution):
                                     0.5 * self.df, self.dimension),
                                 beta=0.5,
                                 dtype=self.dtype,
-                                seed=seed)
+                                seed=distribution_util.gen_new_seed(
+                                    seed, "wishart"))
 
     # Complexity: O(nbk^2)
     x = array_ops.matrix_band_part(x, -1, 0)  # Tri-lower.