commit 6d77b7f3f4c5147b0da526cda2df1946297a6522
author Sagun Bajra <sagunb@google.com> Mon Dec 27 14:13:48 2021 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> Mon Dec 27 14:20:16 2021 -0800
tree 72262a3f0a545893d5d30bf3427876c552c4a1bf
parent b3489fb2d6d5e30ecc1c8b7766f9febd6506b7e6

Internal change

PiperOrigin-RevId: 418528239
Change-Id: I258a6fbfdb98bcf391abf592d5ce0699cb01d900

tensorflow/python/framework/test_util.py

diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index b0755f9..7554aee 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -1130,6 +1130,7 @@
   return wrapper
 
 
+@tf_export("test.with_eager_op_as_function")
 def with_eager_op_as_function(cls=None, only_as_function=False):
   """Adds methods that call original methods with eager_op_as_function enabled.
 

tensorflow/tools/api/golden/v1/tensorflow.test.pbtxt

diff --git a/tensorflow/tools/api/golden/v1/tensorflow.test.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.test.pbtxt
index 106b8ad..5c85c06 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.test.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.test.pbtxt
@@ -76,4 +76,8 @@
     name: "test_src_dir_path"
     argspec: "args=[\'relative_path\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "with_eager_op_as_function"
+    argspec: "args=[\'cls\', \'only_as_function\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+  }
 }

tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt
index 791f16d..382ca5e 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt
@@ -56,4 +56,8 @@
     name: "main"
     argspec: "args=[\'argv\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
+  member_method {
+    name: "with_eager_op_as_function"
+    argspec: "args=[\'cls\', \'only_as_function\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+  }
 }