Test autograph indirect tf.map_fn decorator

tensorflow · May 14, 2020 · 10c7f27 · 10c7f27
1 parent a04c8be
commit 10c7f27
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/tensorflow/python/kernel_tests/map_fn_test.py b/tensorflow/python/kernel_tests/map_fn_test.py
@@ -186,6 +186,24 @@ def testMap_MultiInputSameStructureOutput(self):
     self.assertAllEqual(-nums, received[1])
     self.assertAllEqual(nums, received[2])
 
+  @test_util.run_in_graph_and_eager_modes
+  def testMap_autograph_indirect():
+    def test_function(x):
+      cond = tf.constant(-1)
+      if cond == 0:
+        result = x
+      else:
+        result = x
+      return result
+
+    @tf.function
+    def map_call(x):
+      tf.map_fn(test_function, x)
+
+    x = constant_op.constant([1])
+    y = map_call(x)
+    self.assertAllEqual([1], self.evaluate(y))
+
   @test_util.run_in_graph_and_eager_modes
   def testMapShape(self):
     x = constant_op.constant([[1, 2, 3], [4, 5, 6]])