Move from deprecated self.test_session() to self.cached_session().

self.test_session() has been deprecated in 9962eb5e84b15e309410071b06c2ed2d6148ed44 as its name confuses readers of the test. Moving to cached_session() instead which is more explicit about:
* the fact that the session may be reused.
* the session is not closed even when doing a "with self.test_session()" statement.

PiperOrigin-RevId: 209700671

8 files changed, 73 insertions, 73 deletions

diff --git a/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py b/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py
index 0e65d419a3..76b31a5730 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py
@@ -155,7 +155,7 @@ class EstimatorTest(test.TestCase):
 
   def test_cov_update_thunks(self):
     """Ensures covariance update ops run once per global_step."""
-    with self._graph.as_default(), self.test_session() as sess:
+    with self._graph.as_default(), self.cached_session() as sess:
       fisher_estimator = estimator.FisherEstimatorRoundRobin(
           variables=[self.weights],
           layer_collection=self.layer_collection,
@@ -241,7 +241,7 @@ class EstimatorTest(test.TestCase):
 
   def test_inv_update_thunks(self):
     """Ensures inverse update ops run once per global_step."""
-    with self._graph.as_default(), self.test_session() as sess:
+    with self._graph.as_default(), self.cached_session() as sess:
       fisher_estimator = estimator.FisherEstimatorRoundRobin(
           variables=[self.weights],
           layer_collection=self.layer_collection,
diff --git a/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py b/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py
index 86ec7a095a..f845def507 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py
@@ -59,7 +59,7 @@ def _make_psd(dim):
 class UtilsTest(test.TestCase):
 
   def testComputePiTracenorm(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       diag = ops.convert_to_tensor([1., 2., 0., 1.])
       left_factor = lo.LinearOperatorDiag(diag)
@@ -103,7 +103,7 @@ class FullFBTest(test.TestCase):
       block.instantiate_factors(grads, 0.5)
 
   def testMultiplyInverseTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = (array_ops.constant([1., 2.]), array_ops.constant(3.))
       block = fb.FullFB(lc.LayerCollection(), params)
@@ -124,7 +124,7 @@ class FullFBTest(test.TestCase):
       self.assertAllClose(sess.run(vector * 2 / 3.), sess.run(output))
 
   def testMultiplyInverseNotTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = array_ops.constant([[1.], [2.]])
       block = fb.FullFB(lc.LayerCollection(), params)
@@ -145,7 +145,7 @@ class FullFBTest(test.TestCase):
       self.assertAllClose(sess.run(vector * 2 / 3.), sess.run(output))
 
   def testMultiplyInverseAgainstExplicit(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = (array_ops.constant([1., 2.]), array_ops.constant(3.))
       block = fb.FullFB(lc.LayerCollection(), params)
@@ -203,7 +203,7 @@ class NaiveDiagonalFBTest(test.TestCase):
       block.instantiate_factors(grads, 0.5)
 
   def testMultiplyInverseTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = (array_ops.constant([1., 2.]), array_ops.constant(3.))
       block = fb.NaiveDiagonalFB(lc.LayerCollection(), params)
@@ -222,7 +222,7 @@ class NaiveDiagonalFBTest(test.TestCase):
       self.assertAllClose(sess.run(vector * 2 / 3.), sess.run(output))
 
   def testMultiplyInverseNotTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = array_ops.constant([[1.], [2.]])
       block = fb.NaiveDiagonalFB(lc.LayerCollection(), params)
@@ -240,7 +240,7 @@ class NaiveDiagonalFBTest(test.TestCase):
       self.assertAllClose(sess.run(vector * 2 / 3.), sess.run(output))
 
   def testMultiplyInverseAgainstExplicit(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = (array_ops.constant([1., 2.]), array_ops.constant(3.))
       block = fb.NaiveDiagonalFB(lc.LayerCollection(), params)
@@ -382,7 +382,7 @@ class FullyConnectedDiagonalFBTest(test.TestCase):
       multiply_result: Result of FisherBlock.multiply(params)
       multiply_inverse_result: Result of FisherBlock.multiply_inverse(params)
     """
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       inputs = as_tensors(inputs)
       outputs = as_tensors(outputs)
       output_grads = as_tensors(output_grads)
@@ -425,7 +425,7 @@ class EmbeddingKFACFBTest(test.TestCase):
       block.instantiate_factors(((grads,),), damping)
 
   def testMultiplyInverse(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
 
       # Create a Fisher Block.
@@ -503,7 +503,7 @@ class FullyConnectedKFACBasicFBTest(test.TestCase):
       block.instantiate_factors(((grads,),), 0.5)
 
   def testMultiplyInverseTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       inputs = array_ops.constant([[1., 2., 3.], [3., 4., 5.], [5., 6., 7.]])
       outputs = array_ops.constant([[3., 4.], [5., 6.]])
@@ -535,7 +535,7 @@ class FullyConnectedKFACBasicFBTest(test.TestCase):
       self.assertAllClose([0.343146, 0.686291], output[1])
 
   def testMultiplyInverseNotTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       inputs = array_ops.constant([[1., 2.], [3., 4.]])
       outputs = array_ops.constant([[3., 4.], [5., 6.]])
@@ -561,7 +561,7 @@ class FullyConnectedKFACBasicFBTest(test.TestCase):
                           sess.run(output))
 
   def testMultiplyInverseAgainstExplicit(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       input_dim, output_dim = 3, 2
       inputs = array_ops.zeros([32, input_dim])
@@ -757,7 +757,7 @@ class ConvDiagonalFBTest(test.TestCase):
       multiply_result: Result of FisherBlock.multiply(params)
       multiply_inverse_result: Result of FisherBlock.multiply_inverse(params)
     """
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       inputs = as_tensors(inputs)
       outputs = as_tensors(outputs)
       output_grads = as_tensors(output_grads)
@@ -795,7 +795,7 @@ class DepthwiseConvKFCBasicFBTest(test.TestCase):
       block.instantiate_factors(([grads],), 0.5)
 
   def testMultiplyInverse(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = random_ops.random_normal((3, 3, 8, 2))
       inputs = random_ops.random_normal((32, 5, 5, 8))
@@ -851,7 +851,7 @@ class ConvKFCBasicFBTest(test.TestCase):
     self._testConvKFCBasicFBInitParams([np.ones([1, 2, 2])])
 
   def testMultiplyInverseTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = random_ops.random_normal((2, 2, 2, 2))
       inputs = random_ops.random_normal((2, 2, 2, 2))
@@ -882,7 +882,7 @@ class ConvKFCBasicFBTest(test.TestCase):
       self.assertAllClose([0.27291, 0.409365], output[1])
 
   def testMultiplyInverseNotTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = random_ops.random_normal((2, 2, 2, 2))
       inputs = random_ops.random_normal((2, 2, 2, 2))
@@ -910,7 +910,7 @@ class ConvKFCBasicFBTest(test.TestCase):
       self.assertAllClose([0.136455, 0.27291], sess.run(output)[0])
 
   def testMultiplyInverseNotTupleWithBias(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = [random_ops.random_normal((2, 2, 2, 2))]
       inputs = random_ops.random_normal((2, 2, 2, 2))
@@ -938,7 +938,7 @@ class ConvKFCBasicFBTest(test.TestCase):
       self.assertAllClose([0.136455, 0.27291], sess.run(output)[0])
 
   def testMultiplyInverseAgainstExplicit(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       params = array_ops.zeros((2, 2, 2, 2))
       inputs = array_ops.zeros((2, 2, 2, 2))
diff --git a/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py b/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py
index fad47cd02f..a396ca3f85 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py
@@ -148,7 +148,7 @@ class DenseSquareMatrixFactorTestingDummy(ff.DenseSquareMatrixFactor):
 class NumericalUtilsTest(test.TestCase):
 
   def testComputeCovAgainstNumpy(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       npr.seed(0)
       random_seed.set_random_seed(200)
 
@@ -159,7 +159,7 @@ class NumericalUtilsTest(test.TestCase):
       self.assertAllClose(sess.run(cov), np_cov)
 
   def testComputeCovAgainstNumpyWithAlternativeNormalizer(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       npr.seed(0)
       random_seed.set_random_seed(200)
 
@@ -171,7 +171,7 @@ class NumericalUtilsTest(test.TestCase):
       self.assertAllClose(sess.run(cov), np_cov)
 
   def testAppendHomog(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       npr.seed(0)
 
       m, n = 3, 4
@@ -316,7 +316,7 @@ class DenseSquareMatrixFactorTest(test.TestCase):
       self.assertEqual(0, len(factor.make_inverse_update_ops()))
 
   def testMakeInverseUpdateOpsManyInversesEigenDecomp(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       cov = np.array([[1., 2.], [3., 4.]])
       factor = DenseSquareMatrixFactorTestingDummy(cov.shape)
@@ -348,7 +348,7 @@ class DenseSquareMatrixFactorTest(test.TestCase):
           self.assertNotEqual(new_invs[i][0][0], new_invs[j][0][0])
 
   def testMakeInverseUpdateOpsMatPowerEigenDecomp(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       cov = np.array([[6., 2.], [2., 4.]])
       factor = DenseSquareMatrixFactorTestingDummy(cov.shape)
@@ -369,7 +369,7 @@ class DenseSquareMatrixFactorTest(test.TestCase):
       self.assertAllClose(matpower, matpower_np)
 
   def testMakeInverseUpdateOpsNoEigenDecomp(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       cov = np.array([[5., 2.], [2., 4.]])  # NOTE(mattjj): must be symmetric
       factor = DenseSquareMatrixFactorTestingDummy(cov.shape)
@@ -415,7 +415,7 @@ class FullFactorTest(test.TestCase):
       self.assertEqual([6, 6], cov.get_shape().as_list())
 
   def testMakeCovarianceUpdateOp(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       tensor = array_ops.constant([1., 2.], name='a/b/c')
       factor = ff.FullFactor((tensor,), 2)
@@ -448,7 +448,7 @@ class NaiveDiagonalFactorTest(test.TestCase):
       self.assertEqual([6, 1], cov.get_shape().as_list())
 
   def testMakeCovarianceUpdateOp(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       tensor = array_ops.constant([1., 2.], name='a/b/c')
       factor = ff.NaiveDiagonalFactor((tensor,), 2)
@@ -478,7 +478,7 @@ class EmbeddingInputKroneckerFactorTest(test.TestCase):
       factor.instantiate_cov_variables()
       cov_update_op = factor.make_covariance_update_op(0.0)
 
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         new_cov = sess.run(cov_update_op)
         self.assertAllClose(np.array([1., 1., 0., 0., 1.]) / 3., new_cov)
@@ -555,7 +555,7 @@ class ConvDiagonalFactorTest(test.TestCase):
 
       # Ensure new covariance value is same as outer-product of inputs/outputs
       # vectorized, squared.
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         cov = sess.run(cov_update_op)
         expected_cov = np.outer(inputs.flatten(), outputs_grad.flatten())**2
@@ -591,7 +591,7 @@ class ConvDiagonalFactorTest(test.TestCase):
 
       # Ensure update op doesn't crash.
       cov_update_op = factor.make_covariance_update_op(0.0)
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         sess.run(cov_update_op)
 
@@ -620,7 +620,7 @@ class FullyConnectedKroneckerFactorTest(test.TestCase):
       self._testFullyConnectedKroneckerFactorInit(True, [4, 4], dtype=dtype)
 
   def testMakeCovarianceUpdateOpWithBias(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c')
       factor = ff.FullyConnectedKroneckerFactor(((tensor,),), has_bias=True)
@@ -631,7 +631,7 @@ class FullyConnectedKroneckerFactorTest(test.TestCase):
       self.assertAllClose([[3, 3.5, 1], [3.5, 5.5, 1.5], [1, 1.5, 1]], new_cov)
 
   def testMakeCovarianceUpdateOpNoBias(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c')
       factor = ff.FullyConnectedKroneckerFactor(((tensor,),))
@@ -680,7 +680,7 @@ class ConvInputKroneckerFactorTest(ConvFactorTestCase):
                        factor.get_cov().shape.as_list())
 
       # Ensure cov_update_op doesn't crash.
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         sess.run(factor.make_covariance_update_op(0.0))
         cov = sess.run(factor.get_cov())
@@ -711,7 +711,7 @@ class ConvInputKroneckerFactorTest(ConvFactorTestCase):
                        factor.get_cov().shape.as_list())
 
       # Ensure cov_update_op doesn't crash.
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         sess.run(factor.make_covariance_update_op(0.0))
         cov = sess.run(factor.get_cov())
@@ -736,7 +736,7 @@ class ConvInputKroneckerFactorTest(ConvFactorTestCase):
           has_bias=False)
       factor.instantiate_cov_variables()
 
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         sess.run(factor.make_covariance_update_op(0.0))
         cov = sess.run(factor.get_cov())
@@ -762,7 +762,7 @@ class ConvInputKroneckerFactorTest(ConvFactorTestCase):
           has_bias=False)
       factor.instantiate_cov_variables()
 
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         sess.run(factor.make_covariance_update_op(0.0))
         cov = sess.run(factor.get_cov())
@@ -805,7 +805,7 @@ class ConvInputKroneckerFactorTest(ConvFactorTestCase):
                        cov.get_shape().as_list())
 
   def testMakeCovarianceUpdateOpWithBias(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       input_shape = (2, 1, 1, 1)
       tensor = array_ops.constant(
           np.arange(1, 1 + np.prod(input_shape)).reshape(input_shape).astype(
@@ -824,7 +824,7 @@ class ConvInputKroneckerFactorTest(ConvFactorTestCase):
           new_cov)
 
   def testMakeCovarianceUpdateOpNoBias(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       input_shape = (2, 1, 1, 1)
       tensor = array_ops.constant(
           np.arange(1, 1 + np.prod(input_shape)).reshape(input_shape).astype(
@@ -867,7 +867,7 @@ class ConvOutputKroneckerFactorTest(ConvFactorTestCase):
       ],))
       factor.instantiate_cov_variables()
 
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         sess.run(tf_variables.global_variables_initializer())
         sess.run(factor.make_covariance_update_op(0.0))
         cov = sess.run(factor.get_cov())
@@ -896,7 +896,7 @@ class ConvOutputKroneckerFactorTest(ConvFactorTestCase):
       self.assertEqual([5, 5], cov.get_shape().as_list())
 
   def testMakeCovarianceUpdateOp(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       tensor = np.arange(1, 17).reshape(2, 2, 2, 2).astype(np.float32)
       factor = ff.ConvOutputKroneckerFactor(((array_ops.constant(tensor),),))
@@ -929,7 +929,7 @@ class FullyConnectedMultiKFTest(test.TestCase):
       self.assertEqual([3, 3], cov.get_shape().as_list())
 
   def testMakeCovarianceUpdateOpWithBias(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c')
       factor = ff.FullyConnectedMultiKF(((tensor,),), has_bias=True)
@@ -940,7 +940,7 @@ class FullyConnectedMultiKFTest(test.TestCase):
       self.assertAllClose([[3, 3.5, 1], [3.5, 5.5, 1.5], [1, 1.5, 1]], new_cov)
 
   def testMakeCovarianceUpdateOpNoBias(self):
-    with tf_ops.Graph().as_default(), self.test_session() as sess:
+    with tf_ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c')
       factor = ff.FullyConnectedMultiKF(((tensor,),))
diff --git a/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py b/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py
index cb80fca370..586fcd4c3c 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py
@@ -246,7 +246,7 @@ class LayerCollectionTest(test.TestCase):
     self.assertIn('was already registered', str(cm.exception))
 
   def testRegisterCategoricalPredictiveDistribution(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       logits = linalg_ops.eye(2)
 
@@ -342,7 +342,7 @@ class LayerCollectionTest(test.TestCase):
       lc.register_categorical_predictive_distribution(logits, seed=200)
 
   def testRegisterCategoricalPredictiveDistributionSpecifiedTargets(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       logits = array_ops.constant([[1., 2.], [3., 4.]], dtype=dtypes.float32)
       lc = layer_collection.LayerCollection()
@@ -353,7 +353,7 @@ class LayerCollectionTest(test.TestCase):
       self.assertAlmostEqual(1.6265233, single_loss)
 
   def testRegisterNormalPredictiveDistribution(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       predictions = array_ops.constant(
           [[1., 2.], [3., 4]], dtype=dtypes.float32)
@@ -370,7 +370,7 @@ class LayerCollectionTest(test.TestCase):
       self.assertAlmostEqual(2 * single_loss, double_loss)
 
   def testRegisterNormalPredictiveDistributionSpecifiedTargets(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       predictions = array_ops.constant(
           [[1., 2.], [3., 4.]], dtype=dtypes.float32)
diff --git a/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py b/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py
index c00af5593f..f424e02360 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py
@@ -38,7 +38,7 @@ class InsertSliceInZerosTest(test.TestCase):
     input_tensor = constant_op.constant([[[1, 2]], [[3, 4]]])
     expected_output_array = [[[1, 2], [0, 0]], [[3, 4], [0, 0]]]
     op = loss_functions.insert_slice_in_zeros(input_tensor, 1, 2, 0)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       actual_output_array = sess.run(op)
     self.assertAllEqual(expected_output_array, actual_output_array)
 
@@ -47,7 +47,7 @@ class CategoricalLogitsNegativeLogProbLossTest(test.TestCase):
 
   def testSample(self):
     """Ensure samples can be drawn."""
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.asarray([
           [0., 0., 0.],  #
           [1., -1., 0.]
@@ -60,7 +60,7 @@ class CategoricalLogitsNegativeLogProbLossTest(test.TestCase):
 
   def testEvaluateOnTargets(self):
     """Ensure log probability can be evaluated correctly."""
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.asarray([
           [0., 0., 0.],  #
           [1., -1., 0.]
@@ -83,7 +83,7 @@ class CategoricalLogitsNegativeLogProbLossTest(test.TestCase):
 
   def testEvaluateOnSample(self):
     """Ensure log probability of a sample can be drawn."""
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.asarray([
           [0., 0., 0.],  #
           [1., -1., 0.]
@@ -97,7 +97,7 @@ class CategoricalLogitsNegativeLogProbLossTest(test.TestCase):
       neg_log_prob = sess.run(neg_log_prob)
 
   def testMultiplyFisherSingleVector(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.array([1., 2., 3.])
       loss = loss_functions.CategoricalLogitsNegativeLogProbLoss(logits)
 
@@ -116,7 +116,7 @@ class CategoricalLogitsNegativeLogProbLossTest(test.TestCase):
         self.assertAllClose(expected_result, sess.run(result))
 
   def testMultiplyFisherBatch(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.array([[1., 2., 3.], [4., 6., 8.]])
       loss = loss_functions.CategoricalLogitsNegativeLogProbLoss(logits)
 
@@ -137,7 +137,7 @@ class OnehotCategoricalLogitsNegativeLogProbLossTest(test.TestCase):
 
   def testSample(self):
     """Ensure samples can be drawn."""
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.asarray([
           [0., 0., 0.],  #
           [1., -1., 0.]
@@ -150,7 +150,7 @@ class OnehotCategoricalLogitsNegativeLogProbLossTest(test.TestCase):
 
   def testEvaluateOnTargets(self):
     """Ensure log probability can be evaluated correctly."""
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.asarray([
           [0., 0., 0.],  #
           [1., -1., 0.]
@@ -173,7 +173,7 @@ class OnehotCategoricalLogitsNegativeLogProbLossTest(test.TestCase):
 
   def testEvaluateOnSample(self):
     """Ensure log probability of a sample can be drawn."""
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       logits = np.asarray([
           [0., 0., 0.],  #
           [1., -1., 0.]
diff --git a/tensorflow/contrib/kfac/python/kernel_tests/op_queue_test.py b/tensorflow/contrib/kfac/python/kernel_tests/op_queue_test.py
index b20a70e4ca..4fae4374e1 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/op_queue_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/op_queue_test.py
@@ -36,7 +36,7 @@ class OpQueueTest(test.TestCase):
       ]
       queue = op_queue.OpQueue(ops, seed=0)
 
-      with self.test_session() as sess:
+      with self.cached_session() as sess:
         # Ensure every inv update op gets selected.
         selected_ops = set([queue.next_op(sess) for _ in ops])
         self.assertEqual(set(ops), set(selected_ops))
diff --git a/tensorflow/contrib/kfac/python/kernel_tests/optimizer_test.py b/tensorflow/contrib/kfac/python/kernel_tests/optimizer_test.py
index 560a9b0b42..0b0de12ce6 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/optimizer_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/optimizer_test.py
@@ -84,7 +84,7 @@ class OptimizerTest(test.TestCase):
           momentum_type='regular')
 
   def testSquaredFisherNorm(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       grads_and_vars = [(array_ops.constant([[1., 2.], [3., 4.]]), None),
                         (array_ops.constant([[2., 3.], [4., 5.]]), None)]
       pgrads_and_vars = [(array_ops.constant([[3., 4.], [5., 6.]]), None),
@@ -94,7 +94,7 @@ class OptimizerTest(test.TestCase):
       self.assertAlmostEqual(174., sess.run(sq_norm), places=5)
 
   def testUpdateClipCoeff(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       grads_and_vars = [(array_ops.constant([[1., 2.], [3., 4.]]), None),
                         (array_ops.constant([[2., 3.], [4., 5.]]), None)]
       pgrads_and_vars = [(array_ops.constant([[3., 4.], [5., 6.]]), None),
@@ -129,7 +129,7 @@ class OptimizerTest(test.TestCase):
     pass
 
   def testUpdateVelocities(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       layers = lc.LayerCollection()
       layers.register_categorical_predictive_distribution(
           array_ops.constant([1.0]))
@@ -167,7 +167,7 @@ class OptimizerTest(test.TestCase):
         self.assertFalse(np.equal(first, second).all())
 
   def testApplyGradients(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       layer_collection = lc.LayerCollection()
 
       inputs = array_ops.ones((2, 1)) * 2
diff --git a/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py b/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py
index 2cee01212a..7df79a3c7f 100644
--- a/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py
+++ b/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py
@@ -129,7 +129,7 @@ class UtilsTest(test.TestCase):
     return (weights, biases)
 
   def testFullyConnectedLayerParamsTupleToMat2d(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       layer_params = self._fully_connected_layer_params()
       output = utils.layer_params_to_mat2d(layer_params)
@@ -138,7 +138,7 @@ class UtilsTest(test.TestCase):
           sess.run(output), np.array([[1., 2.], [4., 3.], [1., 2.]]))
 
   def testFullyConnectedLayerParamsTensorToMat2d(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       layer_params = self._fully_connected_layer_params()
       output = utils.layer_params_to_mat2d(layer_params[0])
@@ -153,7 +153,7 @@ class UtilsTest(test.TestCase):
       self.assertListEqual([2 * 2 * 3 + 1, 4], output.get_shape().as_list())
 
   def testKron(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       mat1 = np.array([[1., 2.], [3., 4.]])
       mat2 = np.array([[5., 6.], [7., 8.]])
       mat1_tf = array_ops.constant(mat1)
@@ -163,7 +163,7 @@ class UtilsTest(test.TestCase):
       self.assertAllClose(ans_tf, ans_np)
 
   def testMat2dToFullyConnectedLayerParamsTuple(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       vector_template = self._fully_connected_layer_params()
       mat2d = array_ops.constant([[5., 4.], [3., 2.], [1., 0.]])
@@ -177,7 +177,7 @@ class UtilsTest(test.TestCase):
       self.assertAllClose(b, np.array([1., 0.]))
 
   def testMat2dToFullyConnectedLayerParamsTensor(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       vector_template = self._fully_connected_layer_params()[0]
       mat2d = array_ops.constant([[5., 4.], [3., 2.]])
@@ -187,7 +187,7 @@ class UtilsTest(test.TestCase):
       self.assertAllClose(output, np.array([[5., 4.], [3., 2.]]))
 
   def testTensorsToColumn(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
 
       vector = array_ops.constant(np.array([[0., 1.], [2., 3.]]))
@@ -211,7 +211,7 @@ class UtilsTest(test.TestCase):
           np.array([1., 2., 4., 3., 1., 2., 6., 7., 8., 9.])[:, None])
 
   def testColumnToTensors(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
 
       vector_template = array_ops.constant(np.array([[0., 1.], [2., 3.]]))
@@ -241,7 +241,7 @@ class UtilsTest(test.TestCase):
       self.assertAllClose(c, np.array([[6.], [7.], [8.], [9.]]))
 
   def testPosDefInvCholesky(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       npr.seed(0)
       square = lambda x: np.dot(x, x.T)
@@ -256,7 +256,7 @@ class UtilsTest(test.TestCase):
       self.assertAllClose(sess.run(tf_inv), np_inv)
 
   def testPosDefInvMatrixInverse(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       random_seed.set_random_seed(200)
       npr.seed(0)
       square = lambda x: np.dot(x, x.T)
@@ -296,7 +296,7 @@ class UtilsTest(test.TestCase):
     def increment_var(var):
       return lambda: var.assign_add(1)
 
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       i = variable_scope.get_variable('i', initializer=0)
       accumulators = [
           variable_scope.get_variable('var%d' % j, initializer=0)
@@ -328,7 +328,7 @@ class UtilsTest(test.TestCase):
           values)
 
   def testExtractConvolutionPatches(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       batch_size = 10
       image_spatial_shape = [9, 10, 11]
       in_channels = out_channels = 32
@@ -373,7 +373,7 @@ class UtilsTest(test.TestCase):
       self.assertAllClose(outputs_.flatten(), outputs_flat_.flatten())
 
   def testExtractPointwiseConv2dPatches(self):
-    with ops.Graph().as_default(), self.test_session() as sess:
+    with ops.Graph().as_default(), self.cached_session() as sess:
       batch_size = 10
       image_height = image_width = 8
       in_channels = out_channels = 3