From ded40f9b8f2b99f683d3919a55f93ee4ae5de808 Mon Sep 17 00:00:00 2001
From: Dan Moldovan <mdan@google.com>
Date: Thu, 12 Nov 2020 15:33:24 -0800
Subject: [PATCH] Don't use the heavyweight tf_stack in tf_decorator, which
only needs the name of the parent frame. Remove package dependency.
PiperOrigin-RevId: 342145061
Change-Id: Ife96610c6795c0dac973332e37da1790b58491f0
---
tensorflow/python/BUILD | 4 ++--
tensorflow/python/util/tf_decorator.py | 5 +----
2 files changed, 3 insertions(+), 6 deletions(-)
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 55480f5b95e..16b84895751 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -5696,12 +5696,12 @@ py_strict_library(
"//learning/deepmind/research/language/translation/lm:__subpackages__",
],
deps = [
- ":tf_stack",
"@six_archive//:six",
],
)
-# Leaf library: may not depend on anything else inside TensorFlow.
+# Note: this is a heavyweight library specialized for TensorFlow graphs. Do not use for
+# other purposes.
py_strict_library(
name = "tf_stack",
srcs = ["util/tf_stack.py"],
diff --git a/tensorflow/python/util/tf_decorator.py b/tensorflow/python/util/tf_decorator.py
index 92233f27aea..135811a95df 100644
--- a/tensorflow/python/util/tf_decorator.py
+++ b/tensorflow/python/util/tf_decorator.py
@@ -61,8 +61,6 @@ from __future__ import print_function
import inspect
-from tensorflow.python.util import tf_stack
-
def make_decorator(target,
decorator_func,
@@ -84,8 +82,7 @@ def make_decorator(target,
The `decorator_func` argument with new metadata attached.
"""
if decorator_name is None:
- frame = tf_stack.extract_stack(limit=2)[0]
- decorator_name = frame.name
+ decorator_name = inspect.currentframe().f_back.f_code.co_name
decorator = TFDecorator(decorator_name, target, decorator_doc,
decorator_argspec)
setattr(decorator_func, '_tf_decorator', decorator)