Move from deprecated self.test_session() to self.session() when a graph is set.

self.test_session() has been deprecated in cl/208545396 as its behavior confuses readers of the test. Moving to self.session() instead.

PiperOrigin-RevId: 209696110

1 files changed, 9 insertions, 8 deletions

diff --git a/tensorflow/contrib/compiler/jit_test.py b/tensorflow/contrib/compiler/jit_test.py
index a56a01b163..42b3b9f026 100644
--- a/tensorflow/contrib/compiler/jit_test.py
+++ b/tensorflow/contrib/compiler/jit_test.py
@@ -48,7 +48,7 @@ class JITTest(test.TestCase):
 
   def compute(self, use_jit, compute_fn):
     random_seed.set_random_seed(1234)
-    with self.test_session(graph=ops.Graph()) as sess:
+    with self.session(graph=ops.Graph()) as sess:
       with jit.experimental_jit_scope(use_jit):
         r = compute_fn()
       sess.run(variables.global_variables_initializer())
@@ -88,7 +88,7 @@ class JITTest(test.TestCase):
     self.assertAllClose(v_false_1, v_true_1)
 
   def testJITXlaScope(self):
-    with self.test_session(graph=ops.Graph()):
+    with self.session(graph=ops.Graph()):
       with jit.experimental_jit_scope(True):
         # XlaScope 0
         a1 = constant_op.constant(1)
@@ -138,7 +138,8 @@ class JITTest(test.TestCase):
     self.assertAllClose(v_false_1, v_true_1)
 
   def testDefunNoJitScope(self):
-    with self.test_session(graph=ops.Graph()):
+    with self.session(graph=ops.Graph()):
+
       @function.Defun(compiled=True, noinline=True)
       def mulop(x1, x2):
         return x1 * x2
@@ -153,7 +154,7 @@ class JITTest(test.TestCase):
       self.assertEqual(b"function_mulop", func_attrs["_XlaScope"].s)
 
   def testDefunInheritsJitScope(self):
-    with self.test_session(graph=ops.Graph()):
+    with self.session(graph=ops.Graph()):
       with jit.experimental_jit_scope(True):
         @function.Defun(compiled=True, noinline=True)
         def mulop(x1, x2):
@@ -195,7 +196,7 @@ class CompilationEnabledInGradientTest(test.TestCase):
       self.assertAllClose([[108]], x_grads.eval())
 
   def testCompilationGradientScopeNames(self):
-    with self.test_session(graph=ops.Graph()):
+    with self.session(graph=ops.Graph()):
       with jit.experimental_jit_scope():
         # XlaScope 0
         a1 = constant_op.constant([[1.]])
@@ -217,7 +218,7 @@ class CompilationEnabledInGradientTest(test.TestCase):
       self.assertEqual(b"jit_scope_1", grad_a2.op.get_attr("_XlaScope"))
 
   def testCompilationSeparateGradientScopeNames(self):
-    with self.test_session(graph=ops.Graph()):
+    with self.session(graph=ops.Graph()):
       with jit.experimental_jit_scope(True, separate_compiled_gradients=True):
         # XlaScope 0
         a1 = constant_op.constant([[1.]])
@@ -241,7 +242,7 @@ class CompilationEnabledInGradientTest(test.TestCase):
                        grad_a2.op.get_attr("_XlaScope"))
 
   def testPlaysNicelyWithDefun(self):
-    with self.test_session(graph=ops.Graph()) as sess:
+    with self.session(graph=ops.Graph()) as sess:
       with jit.experimental_jit_scope(True):
         @function.Defun(compiled=True, noinline=True)
         def mulop(x1, x2):
@@ -266,7 +267,7 @@ class CompilationEnabledInGradientTest(test.TestCase):
       self.assertAllClose([1.0, 1.0, 2.0], sess.run([x, r, g_r]))
 
   def testPlaysNicelyWithDefunSeparateGradientScope(self):
-    with self.test_session(graph=ops.Graph()) as sess:
+    with self.session(graph=ops.Graph()) as sess:
       with jit.experimental_jit_scope(True):
 
         @function.Defun(