Autograph decorated functions with tf.map_fn

[bhack]

This is a test to verify the failure of the case described in #26091 (comment)

/cc @dthkao

[bhack]

Mhh.. why test Is passing?

[mdanatg]

I think the issue was fixed a while ago (either in 2.0 or in newer versions, not sure), but the issue ID was never tagged by the fix. I definitely expect the test to pass. Have you seen recent cases when it doesn't work?

[bhack]

We had problems yesterday with 2.2.0 see tensorflow/addons#1830 (comment) and the colab in that ticket.

[bhack]

And also with tf-nightly

[mdanatg]

I see. Do you have a small reproducer in a colab, using @tf.function as decorator? Something simple like this test would do. It might be caused by keyword arguments, if I read through the comments correctly.

[bhack]

def test_function(x):
    cond = tf.constant(-1)

    if cond == 0:
        result = x
    else:
        result = x

    return result

@tf.function
def partial_call(x):
    fn = partial(test_function)
    tf.map_fn(fn, x)

X = tf.random.uniform(shape=(10, 224, 224, 3))
partial_call(X)

[mdanatg]

I don't think that's caused by partial. The problem is that map_fn doesn't call autograph (see below). That's definitely a bug - feel free to file an issue, and if you're interested to help fixing it, I can give a few pointers as well.

A workaround is to wrap the target of map_fn in tf.function:

@tf.function  # Fails without this, but shouldn't.
def test_function(x):
    cond = tf.constant(-1)

    if cond == 0:
        result = x
    else:
        result = x

    return result

@tf.function
def partial_call(x):
    tf.map_fn(test_function, x)

X = tf.random.uniform(shape=(10, 224, 224, 3))
partial_call(X)

[mdanatg]

Let's put this test in https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/kernel_tests/map_fn_test.py, since the fix needs to be applied to map_fn.

[bhack]

@mdanatg I was posting the same if you see the last commit 😸
In the referenced PR yesterday we have solved with the same workaround.
We could try to solve but do you think that it could be a frequent enough case to fix it?

[mdanatg]

It's probably not that frequent, since the functions passed to map_fn are usually simple and don't use control flow. I think it's useful to fix though.

[bhack]

I don't think that we could just use result_value = autograph.tf_convert(fn(elems_value), autograph_ctx.control_status_ctx())

Cause we get:
TypeError: Could not build a TypeSpec for <function convert.<locals>.decorator.<locals>.wrapper at 0x7f651f655e18> with type function

[mdanatg]

I don't think that we could just use result_value = autograph.tf_convert(fn(elems_value), autograph_ctx.control_status_ctx())

Cause we get:
TypeError: Could not build a TypeSpec for <function convert.<locals>.decorator.<locals>.wrapper at 0x7f651f655e18> with type function

Correct, tf_convert doesn't call the function, it just translates it.

So we need something like:

autographed_fn = autograph.tf_convert(fn, autograph_ctx.control_status_ctx())
result_value = autographed_fn(elems_value)

[bhack]

Sorry I needed to restartthe container after a system upgrade and I got:

INFO: Deleting stale sandbox base /root/.cache/bazel/_bazel_root/68a62076e91007a7908bc42a32e4cff9/sandbox

So I will need again hours of recompiling (see tensorflow/build#5 and #39560) 😞

Can you use the Google remote or your local cache to test this a push directly here?

[mdanatg]

I can trigger a build which should finish faster.

[bhack]

Ok I've done a blind commit. If you can start a CI run.

[bhack]

Kokoro we invoke you こころ 🙏 🍖

[bhack]

I don't know if you can disclose this but is the Kokoro waiting list only related to the public queue visible with kokoro:force-run Tests on submitted change or is there an hidden internal one?

[mdanatg]

I'm not sure, actually. Usually it picks up immediately.

[bhack]

import/copybara Pending — Waiting for internal safe review approval is it a manual approval pass right?

[mdanatg]

Yep - looks like the tests passed! Consider updating the PR title.

[bhack]

Yep - looks like the tests passed! Consider updating the PR title.

I hope that the title is clear now.