Fix condition tracing in scale_channel

[bhack]

Try to fix #1829

[bot-of-gabrieldemarmiesse]

@abhichou4

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.

[seanpmorgan]

Thanks! Tracing the encapsulated function appears to work:
https://colab.research.google.com/drive/18QvTLJbS-21QD1tzNyP1G3gjq2SY4RjQ?usp=sharing

Would very much like the standardize the testing for tf.data.Dataset on tfa.image with #992

[bhack]

@seanpmorgan I don't know if in the original colab example, the one with tf.function, there is a know tracing issue with partial.

[WindQAQ]

Many thanks for the quick fix :-) Could we add a simple testcase about tf.data.Dataset? Thanks.

[bhack]

I think that we have a more general issue here about partial with @tf.function.
Instead about tf.data.Dataset tests what is the plan with #992?

[WindQAQ]

I think that we have a more general issue here about partial with @tf.function.
Instead about tf.data.Dataset tests what is the plan with #992?

Could you give me the link to the issue about partial and tf.function? Thank you!

[bhack]

@WindQAQ

fn = tf.function(partial(equalize_image, data_format=data_format))
it works.

@tf_function
def equalize(

doesn't trace equalize_image in partial

@tf_function
def equalize_image(

it works

[bhack]

We had some previous issues with partial:
tensorflow/tensorflow#26091
tensorflow/tensorflow#26602

But seems that they have not tested entering in partial with @tf.function like in your original colab.

[WindQAQ]

Got it! Then I think we could merge this :-) Thanks for the contribution and clarification!

[bhack]

If you are interested I am tying to check if partial inside a decorated function is a bug/limit of the tracer at tensorflow/tensorflow#39541