From 3f8fc30fbb6ae93997765d4fff4f11481c547d5f Mon Sep 17 00:00:00 2001
From: Derek Murray <mrry@google.com>
Date: Tue, 29 Nov 2016 09:03:07 -0800
Subject: [PATCH] Removes `tf.graph_util.set_cpu0()` from the API.

This function was only used internally.
Change: 140485370
---
 tensorflow/python/framework/graph_util.py     |  2 --
 .../python/framework/graph_util_impl.py       | 18 ----------------
 tensorflow/python/training/saver.py           | 21 ++++++++++++++++---
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/tensorflow/python/framework/graph_util.py b/tensorflow/python/framework/graph_util.py
index 402a5ebf0ee1d9..a666630e44b0ec 100644
--- a/tensorflow/python/framework/graph_util.py
+++ b/tensorflow/python/framework/graph_util.py
@@ -26,7 +26,6 @@
 from tensorflow.python.framework.graph_util_impl import extract_sub_graph
 from tensorflow.python.framework.graph_util_impl import must_run_on_cpu
 from tensorflow.python.framework.graph_util_impl import remove_training_nodes
-from tensorflow.python.framework.graph_util_impl import set_cpu0
 from tensorflow.python.framework.graph_util_impl import tensor_shape_from_node_def_name
 # pylint: enable=unused-import
 from tensorflow.python.util.all_util import remove_undocumented
@@ -36,7 +35,6 @@
     "convert_variables_to_constants",
     "extract_sub_graph",
     "must_run_on_cpu",
-    "set_cpu0",
     "tensor_shape_from_node_def_name",
     "remove_training_nodes",
 ]
diff --git a/tensorflow/python/framework/graph_util_impl.py b/tensorflow/python/framework/graph_util_impl.py
index ba693503c259df..587f883260ce3d 100644
--- a/tensorflow/python/framework/graph_util_impl.py
+++ b/tensorflow/python/framework/graph_util_impl.py
@@ -25,7 +25,6 @@
 from tensorflow.core.framework import attr_value_pb2
 from tensorflow.core.framework import graph_pb2
 from tensorflow.core.framework import node_def_pb2
-from tensorflow.python.framework import device as pydev
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
@@ -49,23 +48,6 @@ def _is_variable_op(op):
   return op in _VARIABLE_OPS
 
 
-def set_cpu0(device_string):
-  """Creates a new device string based on `device_string' but using /CPU:0.
-
-   If the device is already on /CPU:0, this is a no-op.
-
-   Args:
-     device_string: A device string.
-
-   Returns:
-     A device string.
-  """
-  parsed_device = pydev.DeviceSpec.from_string(device_string)
-  parsed_device.device_type = "CPU"
-  parsed_device.device_index = 0
-  return parsed_device.to_string()
-
-
 def must_run_on_cpu(node, pin_variables_on_cpu=False):
   """Returns True if the given node_def must run on CPU, otherwise False.
 
diff --git a/tensorflow/python/training/saver.py b/tensorflow/python/training/saver.py
index 68e4bfd0f8743d..cb4e1de2353b7f 100644
--- a/tensorflow/python/training/saver.py
+++ b/tensorflow/python/training/saver.py
@@ -36,7 +36,6 @@
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import device as pydev
 from tensorflow.python.framework import errors
-from tensorflow.python.framework import graph_util
 from tensorflow.python.framework import meta_graph
 from tensorflow.python.framework import ops
 from tensorflow.python.lib.io import file_io
@@ -62,6 +61,23 @@
                      "ResourceGather"])
 
 
+def _set_cpu0(device_string):
+  """Creates a new device string based on `device_string` but using /CPU:0.
+
+  If the device is already on /CPU:0, this is a no-op.
+
+  Args:
+    device_string: A device string.
+
+  Returns:
+    A device string.
+  """
+  parsed_device = pydev.DeviceSpec.from_string(device_string)
+  parsed_device.device_type = "CPU"
+  parsed_device.device_index = 0
+  return parsed_device.to_string()
+
+
 class BaseSaverBuilder(object):
   """Base class for Savers.
 
@@ -380,8 +396,7 @@ def _AddRestoreOps(self,
       # available on the GPU.
       # TODO(touts): Re-enable restore on GPU when we can support annotating
       # string tensors as "HostMemory" inputs.
-      with ops.device(
-          graph_util.set_cpu0(saveable.device) if saveable.device else None):
+      with ops.device(_set_cpu0(saveable.device) if saveable.device else None):
         with ops.control_dependencies(restore_control_inputs):
           tensors = self.restore_op(filename_tensor, saveable, preferred_shard)
           shapes = None