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

1 files changed, 35 insertions, 10 deletions

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,
+])