Remove hourglass imports from even more tests

Change: 142318245

1 files changed, 309 insertions, 254 deletions

diff --git a/tensorflow/python/training/session_manager_test.py b/tensorflow/python/training/session_manager_test.py
index 1d3dcda119..bb6fd84015 100644
--- a/tensorflow/python/training/session_manager_test.py
+++ b/tensorflow/python/training/session_manager_test.py
@@ -12,45 +12,58 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
 """Tests for SessionManager."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
 import os
 
-import tensorflow as tf
-
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
 from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training import server_lib
+from tensorflow.python.training import session_manager
 
 
-class SessionManagerTest(tf.test.TestCase):
+class SessionManagerTest(test.TestCase):
 
   def testPrepareSessionSucceeds(self):
-    with tf.Graph().as_default():
-      v = tf.Variable([1.0, 2.0, 3.0], name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      sess = sm.prepare_session("", init_op=tf.global_variables_initializer())
+    with ops.Graph().as_default():
+      v = variables.Variable([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session(
+          "", init_op=variables.global_variables_initializer())
       self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
 
   def testPrepareSessionSucceedsWithInitFeedDict(self):
-    with tf.Graph().as_default():
-      p = tf.placeholder(tf.float32, shape=(3,))
-      v = tf.Variable(p, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      sess = sm.prepare_session("",
-                                init_op=tf.global_variables_initializer(),
-                                init_feed_dict={p: [1.0, 2.0, 3.0]})
+    with ops.Graph().as_default():
+      p = array_ops.placeholder(dtypes.float32, shape=(3,))
+      v = variables.Variable(p, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          init_feed_dict={p: [1.0, 2.0, 3.0]})
       self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
 
   def testPrepareSessionSucceedsWithInitFn(self):
-    with tf.Graph().as_default():
-      v = tf.Variable([125], name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      sess = sm.prepare_session("",
-                                init_fn=lambda sess: sess.run(v.initializer))
+    with ops.Graph().as_default():
+      v = variables.Variable([125], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session(
+          "", init_fn=lambda sess: sess.run(v.initializer))
       self.assertAllClose([125], sess.run(v))
 
   def testPrepareSessionFails(self):
@@ -60,44 +73,58 @@ class SessionManagerTest(tf.test.TestCase):
       gfile.DeleteRecursively(checkpoint_dir)
       gfile.DeleteRecursively(checkpoint_dir2)
     except errors.OpError:
-      pass                      # Ignore
+      pass  # Ignore
     gfile.MakeDirs(checkpoint_dir)
 
-    with tf.Graph().as_default():
-      v = tf.Variable([1.0, 2.0, 3.0], name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      saver = tf.train.Saver({"v": v})
-      sess = sm.prepare_session("", init_op=tf.global_variables_initializer(),
-                                saver=saver, checkpoint_dir=checkpoint_dir)
+    with ops.Graph().as_default():
+      v = variables.Variable([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          saver=saver,
+          checkpoint_dir=checkpoint_dir)
       self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
       checkpoint_filename = os.path.join(checkpoint_dir,
                                          "prepare_session_checkpoint")
       saver.save(sess, checkpoint_filename)
     # Create a new Graph and SessionManager and recover.
-    with tf.Graph().as_default():
+    with ops.Graph().as_default():
       # Renames the checkpoint directory.
       os.rename(checkpoint_dir, checkpoint_dir2)
       gfile.MakeDirs(checkpoint_dir)
-      v = tf.Variable([6.0, 7.0, 8.0], name="v")
+      v = variables.Variable([6.0, 7.0, 8.0], name="v")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-      tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      saver = tf.train.Saver({"v": v})
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+      session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
       # This should fail as there's no checkpoint within 2 seconds.
       with self.assertRaisesRegexp(
           RuntimeError, "no init_op or init_fn or local_init_op was given"):
-        sess = sm.prepare_session("", init_op=None, saver=saver,
-                                  checkpoint_dir=checkpoint_dir,
-                                  wait_for_checkpoint=True, max_wait_secs=2)
+        sess = sm.prepare_session(
+            "",
+            init_op=None,
+            saver=saver,
+            checkpoint_dir=checkpoint_dir,
+            wait_for_checkpoint=True,
+            max_wait_secs=2)
       # Rename the checkpoint directory back.
       gfile.DeleteRecursively(checkpoint_dir)
       os.rename(checkpoint_dir2, checkpoint_dir)
       # This should succeed as there's checkpoint.
-      sess = sm.prepare_session("", init_op=None, saver=saver,
-                                checkpoint_dir=checkpoint_dir,
-                                wait_for_checkpoint=True, max_wait_secs=2)
+      sess = sm.prepare_session(
+          "",
+          init_op=None,
+          saver=saver,
+          checkpoint_dir=checkpoint_dir,
+          wait_for_checkpoint=True,
+          max_wait_secs=2)
       self.assertEqual(
-          True, tf.is_variable_initialized(
+          True,
+          variables.is_variable_initialized(
               sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
 
   def testRecoverSession(self):
@@ -106,41 +133,44 @@ class SessionManagerTest(tf.test.TestCase):
     try:
       gfile.DeleteRecursively(checkpoint_dir)
     except errors.OpError:
-      pass                      # Ignore
+      pass  # Ignore
     gfile.MakeDirs(checkpoint_dir)
 
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      saver = tf.train.Saver({"v": v})
-      sess, initialized = sm.recover_session("", saver=saver,
-                                             checkpoint_dir=checkpoint_dir)
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertFalse(initialized)
       sess.run(v.initializer)
       self.assertEquals(1, sess.run(v))
-      saver.save(sess, os.path.join(checkpoint_dir,
-                                    "recover_session_checkpoint"))
+      saver.save(sess,
+                 os.path.join(checkpoint_dir, "recover_session_checkpoint"))
     # Create a new Graph and SessionManager and recover.
-    with tf.Graph().as_default():
-      v = tf.Variable(2, name="v")
+    with ops.Graph().as_default():
+      v = variables.Variable(2, name="v")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables())
-      saver = tf.train.Saver({"v": v})
-      sess, initialized = sm2.recover_session("", saver=saver,
-                                              checkpoint_dir=checkpoint_dir)
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm2.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertTrue(initialized)
       self.assertEqual(
-          True, tf.is_variable_initialized(
+          True,
+          variables.is_variable_initialized(
               sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
       self.assertEquals(1, sess.run(v))
 
   def testWaitForSessionReturnsNoneAfterTimeout(self):
-    with tf.Graph().as_default():
-      tf.Variable(1, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables(),
-                                   recovery_wait_secs=1)
+    with ops.Graph().as_default():
+      variables.Variable(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          recovery_wait_secs=1)
 
       # Set max_wait_secs to allow us to try a few times.
       with self.assertRaises(errors.DeadlineExceededError):
@@ -152,9 +182,9 @@ class SessionManagerTest(tf.test.TestCase):
     with self.assertRaisesRegexp(ValueError,
                                  "If you pass a ready_for_local_init_op "
                                  "you must also pass a local_init_op "):
-      tf.train.SessionManager(
-          ready_for_local_init_op=tf.report_uninitialized_variables(
-              tf.global_variables()),
+      session_manager.SessionManager(
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
           local_init_op=None)
 
   def testRecoverSessionWithReadyForLocalInitOp(self):
@@ -167,45 +197,46 @@ class SessionManagerTest(tf.test.TestCase):
       pass  # Ignore
     gfile.MakeDirs(checkpoint_dir)
 
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      saver = tf.train.Saver({"v": v})
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
       sess, initialized = sm.recover_session(
           "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertFalse(initialized)
       sess.run(v.initializer)
       self.assertEquals(1, sess.run(v))
-      saver.save(sess, os.path.join(checkpoint_dir,
-                                    "recover_session_checkpoint"))
+      saver.save(sess,
+                 os.path.join(checkpoint_dir, "recover_session_checkpoint"))
     # Create a new Graph and SessionManager and recover.
-    with tf.Graph().as_default():
-      v = tf.Variable(2, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      v = variables.Variable(2, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables(),
-          ready_for_local_init_op=tf.report_uninitialized_variables(
-              tf.global_variables()),
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
           local_init_op=w.initializer)
-      saver = tf.train.Saver({"v": v})
+      saver = saver_lib.Saver({"v": v})
       sess, initialized = sm2.recover_session(
           "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertTrue(initialized)
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
       self.assertEquals(1, sess.run(v))
       self.assertEquals(1, sess.run(w))
 
@@ -224,44 +255,45 @@ class SessionManagerTest(tf.test.TestCase):
       pass  # Ignore
     gfile.MakeDirs(checkpoint_dir)
 
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
-      saver = tf.train.Saver({"v": v})
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
       sess, initialized = sm.recover_session(
           "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertFalse(initialized)
       sess.run(v.initializer)
       self.assertEquals(1, sess.run(v))
-      saver.save(sess, os.path.join(checkpoint_dir,
-                                    "recover_session_checkpoint"))
+      saver.save(sess,
+                 os.path.join(checkpoint_dir, "recover_session_checkpoint"))
     # Create a new Graph and SessionManager and recover.
-    with tf.Graph().as_default():
-      v = tf.Variable(2, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      v = variables.Variable(2, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables(),
-          ready_for_local_init_op=tf.report_uninitialized_variables(),
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(),
           local_init_op=w.initializer)
-      saver = tf.train.Saver({"v": v})
+      saver = saver_lib.Saver({"v": v})
       sess, initialized = sm2.recover_session(
           "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertFalse(initialized)
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
       self.assertEqual(
           False,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
       self.assertEquals(1, sess.run(v))
 
   def testRecoverSessionNoChkptStillRunsLocalInitOp(self):
@@ -269,16 +301,16 @@ class SessionManagerTest(tf.test.TestCase):
     # In particular, we continue to ensure that recover_session will execute
     # local_init_op exactly once, regardless of whether the session was
     # successfully recovered.
-    with tf.Graph().as_default():
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      w = variables.Variable(
           1,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables(),
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
           ready_for_local_init_op=None,
           local_init_op=w.initializer)
       # Try to recover session from None
@@ -288,8 +320,8 @@ class SessionManagerTest(tf.test.TestCase):
       self.assertFalse(initialized)
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
       self.assertEquals(1, sess.run(w))
 
   def testRecoverSessionFailsStillRunsLocalInitOp(self):
@@ -304,21 +336,21 @@ class SessionManagerTest(tf.test.TestCase):
     gfile.MakeDirs(checkpoint_dir)
 
     # Create a new Graph and SessionManager and recover.
-    with tf.Graph().as_default():
-      v = tf.Variable(2, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      v = variables.Variable(2, name="v")
+      w = variables.Variable(
           1,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables(),
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
           ready_for_local_init_op=None,
           local_init_op=w.initializer)
-      saver = tf.train.Saver({"v": v})
+      saver = saver_lib.Saver({"v": v})
       sess, initialized = sm2.recover_session(
           "",
           saver=saver,
@@ -327,142 +359,142 @@ class SessionManagerTest(tf.test.TestCase):
       self.assertFalse(initialized)
       self.assertEqual(
           False,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
       self.assertEquals(1, sess.run(w))
 
   def testWaitForSessionLocalInit(self):
-    server = tf.train.Server.create_local_server()
-    with tf.Graph().as_default() as graph:
-      v = tf.Variable(1, name="v")
-      w = tf.Variable(
+    server = server_lib.Server.create_local_server()
+    with ops.Graph().as_default() as graph:
+      v = variables.Variable(1, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
-      sm = tf.train.SessionManager(
+      sm = session_manager.SessionManager(
           graph=graph,
-          ready_op=tf.report_uninitialized_variables(),
-          ready_for_local_init_op=tf.report_uninitialized_variables(
-              tf.global_variables()),
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
           local_init_op=w.initializer)
 
       # Initialize v but not w
-      s = tf.Session(server.target, graph=graph)
+      s = session_lib.Session(server.target, graph=graph)
       s.run(v.initializer)
 
       sess = sm.wait_for_session(server.target, max_wait_secs=3)
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
       self.assertEquals(1, sess.run(v))
       self.assertEquals(1, sess.run(w))
 
   def testWaitForSessionWithReadyForLocalInitOpFailsToReadyLocal(self):
-    with tf.Graph().as_default() as graph:
-      v = tf.Variable(1, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default() as graph:
+      v = variables.Variable(1, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
-      sm = tf.train.SessionManager(
+      sm = session_manager.SessionManager(
           graph=graph,
-          ready_op=tf.report_uninitialized_variables(),
-          ready_for_local_init_op=tf.report_uninitialized_variables(),
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(),
           local_init_op=w.initializer)
 
-      with self.assertRaises(tf.errors.DeadlineExceededError):
+      with self.assertRaises(errors_impl.DeadlineExceededError):
         # Time-out because w fails to be initialized,
         # because of overly restrictive ready_for_local_init_op
         sm.wait_for_session("", max_wait_secs=3)
 
   def testWaitForSessionInsufficientReadyForLocalInitCheck(self):
-    with tf.Graph().as_default() as graph:
-      v = tf.Variable(1, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default() as graph:
+      v = variables.Variable(1, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
-      sm = tf.train.SessionManager(
+      sm = session_manager.SessionManager(
           graph=graph,
-          ready_op=tf.report_uninitialized_variables(),
+          ready_op=variables.report_uninitialized_variables(),
           ready_for_local_init_op=None,
           local_init_op=w.initializer)
-    with self.assertRaisesRegexp(tf.errors.FailedPreconditionError,
+    with self.assertRaisesRegexp(errors_impl.FailedPreconditionError,
                                  "Attempting to use uninitialized value v"):
       sm.wait_for_session("", max_wait_secs=3)
 
   def testPrepareSessionWithReadyForLocalInitOp(self):
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables(),
-          ready_for_local_init_op=tf.report_uninitialized_variables(
-              tf.global_variables()),
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
           local_init_op=w.initializer)
       sess = sm2.prepare_session("", init_op=v.initializer)
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
       self.assertEqual(
           True,
-          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
-              session=sess))
+          variables.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
       self.assertEquals(1, sess.run(v))
       self.assertEquals(1, sess.run(w))
 
   def testPrepareSessionDidNotInitLocalVariable(self):
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables())
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
       with self.assertRaisesRegexp(RuntimeError,
                                    "Init operations did not make model ready"):
         sm2.prepare_session("", init_op=v.initializer)
 
   def testPrepareSessionWithReadyNotReadyForLocal(self):
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables(),
-          ready_for_local_init_op=tf.report_uninitialized_variables(
-              tf.global_variables()),
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
           local_init_op=w.initializer)
       with self.assertRaisesRegexp(
           RuntimeError,
@@ -470,50 +502,55 @@ class SessionManagerTest(tf.test.TestCase):
         sm2.prepare_session("", init_op=None)
 
   def testPrepareSessionWithInsufficientReadyForLocalInitCheck(self):
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      w = tf.Variable(
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      w = variables.Variable(
           v,
           trainable=False,
-          collections=[tf.GraphKeys.LOCAL_VARIABLES],
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
           name="w")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-        self.assertEqual(False, tf.is_variable_initialized(w).eval())
-      sm2 = tf.train.SessionManager(
-          ready_op=tf.report_uninitialized_variables(),
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+        self.assertEqual(False, variables.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
           ready_for_local_init_op=None,
           local_init_op=w.initializer)
-    with self.assertRaisesRegexp(tf.errors.FailedPreconditionError,
+    with self.assertRaisesRegexp(errors_impl.FailedPreconditionError,
                                  "Attempting to use uninitialized value v"):
       sm2.prepare_session("", init_op=None)
 
 
-class ObsoleteSessionManagerTest(tf.test.TestCase):
+class ObsoleteSessionManagerTest(test.TestCase):
 
   def testPrepareSessionSucceeds(self):
-    with tf.Graph().as_default():
-      v = tf.Variable([1.0, 2.0, 3.0], name="v")
-      sm = tf.train.SessionManager(ready_op=tf.assert_variables_initialized())
-      sess = sm.prepare_session("", init_op=tf.global_variables_initializer())
+    with ops.Graph().as_default():
+      v = variables.Variable([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      sess = sm.prepare_session(
+          "", init_op=variables.global_variables_initializer())
       self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
 
   def testPrepareSessionSucceedsWithInitFeedDict(self):
-    with tf.Graph().as_default():
-      p = tf.placeholder(tf.float32, shape=(3,))
-      v = tf.Variable(p, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.assert_variables_initialized())
-      sess = sm.prepare_session("",
-                                init_op=tf.global_variables_initializer(),
-                                init_feed_dict={p: [1.0, 2.0, 3.0]})
+    with ops.Graph().as_default():
+      p = array_ops.placeholder(dtypes.float32, shape=(3,))
+      v = variables.Variable(p, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          init_feed_dict={p: [1.0, 2.0, 3.0]})
       self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
 
   def testPrepareSessionSucceedsWithInitFn(self):
-    with tf.Graph().as_default():
-      v = tf.Variable([125], name="v")
-      sm = tf.train.SessionManager(ready_op=tf.assert_variables_initialized())
-      sess = sm.prepare_session("",
-                                init_fn=lambda sess: sess.run(v.initializer))
+    with ops.Graph().as_default():
+      v = variables.Variable([125], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      sess = sm.prepare_session(
+          "", init_fn=lambda sess: sess.run(v.initializer))
       self.assertAllClose([125], sess.run(v))
 
   def testPrepareSessionFails(self):
@@ -523,44 +560,58 @@ class ObsoleteSessionManagerTest(tf.test.TestCase):
       gfile.DeleteRecursively(checkpoint_dir)
       gfile.DeleteRecursively(checkpoint_dir2)
     except errors.OpError:
-      pass                      # Ignore
+      pass  # Ignore
     gfile.MakeDirs(checkpoint_dir)
 
-    with tf.Graph().as_default():
-      v = tf.Variable([1.0, 2.0, 3.0], name="v")
-      sm = tf.train.SessionManager(ready_op=tf.assert_variables_initialized())
-      saver = tf.train.Saver({"v": v})
-      sess = sm.prepare_session("", init_op=tf.global_variables_initializer(),
-                                saver=saver, checkpoint_dir=checkpoint_dir)
+    with ops.Graph().as_default():
+      v = variables.Variable([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          saver=saver,
+          checkpoint_dir=checkpoint_dir)
       self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
       checkpoint_filename = os.path.join(checkpoint_dir,
                                          "prepare_session_checkpoint")
       saver.save(sess, checkpoint_filename)
     # Create a new Graph and SessionManager and recover.
-    with tf.Graph().as_default():
+    with ops.Graph().as_default():
       # Renames the checkpoint directory.
       os.rename(checkpoint_dir, checkpoint_dir2)
       gfile.MakeDirs(checkpoint_dir)
-      v = tf.Variable([6.0, 7.0, 8.0], name="v")
+      v = variables.Variable([6.0, 7.0, 8.0], name="v")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-      tf.train.SessionManager(ready_op=tf.assert_variables_initialized())
-      saver = tf.train.Saver({"v": v})
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+      session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
       # This should fail as there's no checkpoint within 2 seconds.
       with self.assertRaisesRegexp(
           RuntimeError, "no init_op or init_fn or local_init_op was given"):
-        sess = sm.prepare_session("", init_op=None, saver=saver,
-                                  checkpoint_dir=checkpoint_dir,
-                                  wait_for_checkpoint=True, max_wait_secs=2)
+        sess = sm.prepare_session(
+            "",
+            init_op=None,
+            saver=saver,
+            checkpoint_dir=checkpoint_dir,
+            wait_for_checkpoint=True,
+            max_wait_secs=2)
       # Rename the checkpoint directory back.
       gfile.DeleteRecursively(checkpoint_dir)
       os.rename(checkpoint_dir2, checkpoint_dir)
       # This should succeed as there's checkpoint.
-      sess = sm.prepare_session("", init_op=None, saver=saver,
-                                checkpoint_dir=checkpoint_dir,
-                                wait_for_checkpoint=True, max_wait_secs=2)
+      sess = sm.prepare_session(
+          "",
+          init_op=None,
+          saver=saver,
+          checkpoint_dir=checkpoint_dir,
+          wait_for_checkpoint=True,
+          max_wait_secs=2)
       self.assertEqual(
-          True, tf.is_variable_initialized(
+          True,
+          variables.is_variable_initialized(
               sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
 
   def testRecoverSession(self):
@@ -569,40 +620,44 @@ class ObsoleteSessionManagerTest(tf.test.TestCase):
     try:
       gfile.DeleteRecursively(checkpoint_dir)
     except errors.OpError:
-      pass                      # Ignore
+      pass  # Ignore
     gfile.MakeDirs(checkpoint_dir)
 
-    with tf.Graph().as_default():
-      v = tf.Variable(1, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.assert_variables_initialized())
-      saver = tf.train.Saver({"v": v})
-      sess, initialized = sm.recover_session("", saver=saver,
-                                             checkpoint_dir=checkpoint_dir)
+    with ops.Graph().as_default():
+      v = variables.Variable(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertFalse(initialized)
       sess.run(v.initializer)
       self.assertEquals(1, sess.run(v))
-      saver.save(sess, os.path.join(checkpoint_dir,
-                                    "recover_session_checkpoint"))
+      saver.save(sess,
+                 os.path.join(checkpoint_dir, "recover_session_checkpoint"))
     # Create a new Graph and SessionManager and recover.
-    with tf.Graph().as_default():
-      v = tf.Variable(2, name="v")
+    with ops.Graph().as_default():
+      v = variables.Variable(2, name="v")
       with self.test_session():
-        self.assertEqual(False, tf.is_variable_initialized(v).eval())
-      sm2 = tf.train.SessionManager(ready_op=tf.assert_variables_initialized())
-      saver = tf.train.Saver({"v": v})
-      sess, initialized = sm2.recover_session("", saver=saver,
-                                              checkpoint_dir=checkpoint_dir)
+        self.assertEqual(False, variables.is_variable_initialized(v).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm2.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
       self.assertTrue(initialized)
       self.assertEqual(
-          True, tf.is_variable_initialized(
+          True,
+          variables.is_variable_initialized(
               sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
       self.assertEquals(1, sess.run(v))
 
   def testWaitForSessionReturnsNoneAfterTimeout(self):
-    with tf.Graph().as_default():
-      tf.Variable(1, name="v")
-      sm = tf.train.SessionManager(ready_op=tf.assert_variables_initialized(),
-                                   recovery_wait_secs=1)
+    with ops.Graph().as_default():
+      variables.Variable(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized(),
+          recovery_wait_secs=1)
 
       # Set max_wait_secs to allow us to try a few times.
       with self.assertRaises(errors.DeadlineExceededError):
@@ -610,4 +665,4 @@ class ObsoleteSessionManagerTest(tf.test.TestCase):
 
 
 if __name__ == "__main__":
-  tf.test.main()
+  test.main()