From 154df32a959df74b3a1c377ff72f955d755b3d34 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 6 Jul 2017 19:07:18 -0700
Subject: Follow up changes of tfprof migration: 1. Add a option_builder module
 to build options for Python API. 2. Expose necessary profiler protos. 3.
 Rename command line tool from tfprof to profiler.

PiperOrigin-RevId: 161160274
---
 tensorflow/python/profiler/profiler.py | 45 ++++++++++++++++++++++++++--------
 1 file changed, 35 insertions(+), 10 deletions(-)

(limited to 'tensorflow/python/profiler/profiler.py')

diff --git a/tensorflow/python/profiler/profiler.py b/tensorflow/python/profiler/profiler.py
index 19b253ab2f..130dcb5134 100644
--- a/tensorflow/python/profiler/profiler.py
+++ b/tensorflow/python/profiler/profiler.py
@@ -12,21 +12,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""tfprof is a tool that profile various aspect of TensorFlow model.
-
-@@Profiler
-@@profile
-@@advise
-@@write_op_log
+"""profiler python module provides APIs to profile TensorFlow models.
 """
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-# pylint: disable=unused-import,wildcard-import
-from tensorflow.python.profiler.model_analyzer import *
-from tensorflow.python.profiler.tfprof_logger import *
+# pylint: disable=unused-import
+from tensorflow.core.profiler.tfprof_log_pb2 import OpLogProto
+from tensorflow.core.profiler.tfprof_output_pb2 import AdviceProto
+from tensorflow.core.profiler.tfprof_output_pb2 import GraphNodeProto
+from tensorflow.core.profiler.tfprof_output_pb2 import MultiGraphNodeProto
+
+from tensorflow.python.profiler.model_analyzer import advise
+from tensorflow.python.profiler.model_analyzer import profile
+from tensorflow.python.profiler.model_analyzer import Profiler
+from tensorflow.python.profiler.option_builder import ProfileOptionBuilder
+from tensorflow.python.profiler.tfprof_logger import write_op_log
+
 from tensorflow.python.util.all_util import remove_undocumented
 
 
-remove_undocumented(__name__, [])
+_allowed_symbols = [
+    'Profiler',
+    'profile',
+    'ProfileOptionBuilder',
+    'advise',
+    'write_op_log',
+]
+
+_allowed_symbols.extend([
+    'GraphNodeProto',
+    'MultiGraphNodeProto',
+    'AdviceProto',
+    'OpLogProto',
+])
+
+remove_undocumented(__name__, _allowed_symbols, [
+    Profiler,
+    profile,
+    ProfileOptionBuilder,
+    advise,
+    write_op_log,
+])
-- 
cgit v1.2.3