Internal change

PiperOrigin-RevId: 403113437 Change-Id: I79a81903cfaa4817eb9d11998835c177df0e9d8d
Marbck · Oct 14, 2021 · eaae7ec · eaae7ec
1 parent 7c44e35
commit eaae7ec
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tensorflow/python/eager/run_eager_op_as_function_test.py b/tensorflow/python/eager/run_eager_op_as_function_test.py
@@ -48,7 +48,7 @@ def run_benchmark(func, num_iters, unused_execution_mode):
 
 
 # TODO(srbs): Why can't we use absl parameterized here?
-@test_util.with_eager_op_as_function(only_as_function=True)
+@test_util.with_eager_op_as_function
 class MicroBenchmarks(benchmarks_test_base.MicroBenchmarksBase):
 
   def __init__(self):
@@ -78,6 +78,8 @@ def _get_benchmark_name(self):
       raise ValueError("Unable to determine calling Benchmark function.")
     if context.is_tfrt_enabled():
       name = name + "_tfrt"
+    if context.run_eager_op_as_function_enabled():
+      name = name + "_eager_op_as_function"
     return name
 
   def _run(self, func, num_iters):
@@ -90,7 +92,7 @@ def _benchmark_matmul(self, mat, device):
       if device == GPU:
         mat = mat.gpu()
       func = lambda: math_ops.matmul(mat, mat)
-      self._run(func, num_iters=1000)
+      self._run(func, num_iters=5000)
 
   def benchmark_tf_matmul_2_by_2_CPU(self):
     self._benchmark_matmul(self._m_2_by_2, CPU)
@@ -111,7 +113,7 @@ def benchmark_tf_matmul_1000_by_1000_GPU(self):
     self._benchmark_matmul(self._m_1000_by_1000, GPU)
 
 
-@test_util.with_eager_op_as_function(only_as_function=True)
+@test_util.with_eager_op_as_function
 class RunEagerOpAsFunctionTest(test.TestCase):
 
   def setUp(self):