Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Autograph decorated functions with tf.map_fn #39541

Merged
merged 10 commits into from
May 18, 2020
Merged

Autograph decorated functions with tf.map_fn #39541

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
10c7f27
Test autograph indirect tf.map_fn decorator
bhack May 14, 2020
38e941d
Fix missing return
bhack May 14, 2020
ffef546
Fix missing self
bhack May 14, 2020
9a6a647
Add test and remove decorator
bhack May 15, 2020
86342e2
restore a remove export
bhack May 15, 2020
103bb01
Verifiy differences with test annotation
bhack May 15, 2020
bbc2f3a
Let test to fail
bhack May 15, 2020
560762e
Test autograph transform of fn
bhack May 15, 2020
d6dd56f
Remove original fn call
bhack May 15, 2020
64d839b
Fix lint and improve readibility
bhack May 15, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
19 changes: 19 additions & 0 deletions tensorflow/python/kernel_tests/map_fn_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

import numpy as np

from tensorflow.python.eager import def_function
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
Expand Down Expand Up @@ -186,6 +187,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(self):
def test_function(x):
cond = constant_op.constant(-1)
if cond == 0:
result = x
else:
result = x
return result

@def_function.function
def map_call(x):
return map_fn.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]])
Expand Down
6 changes: 5 additions & 1 deletion tensorflow/python/ops/map_fn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@

import re

from tensorflow.python.autograph.core import ag_ctx as autograph_ctx
from tensorflow.python.autograph.impl import api as autograph
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
Expand Down Expand Up @@ -477,7 +479,9 @@ def compute(i, tas):
elems_value_flat = _elems_value_batchable_to_flat(elems_value_batchable,
elems_flat_signature)
elems_value = elems_unflatten(elems_value_flat)
result_value = fn(elems_value)
ag_ctx = autograph_ctx.control_status_ctx()
autographed_fn = autograph.tf_convert(fn, ag_ctx)
result_value = autographed_fn(elems_value)
nest.assert_same_structure(fn_output_signature or elems, result_value)
result_value_flat = nest.flatten(result_value)
result_value_batchable = _result_value_flat_to_batchable(
Expand Down