tf.function condition with tensor

[bhack]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
  Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
  Colab
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary):
  colab
• TensorFlow version (use command below):
  default colab
• Python version:
  default colab
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:
• GPU model and memory:

You can collect some of this information using our environment capture
script
You can also obtain the TensorFlow version with:

1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"
2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior

@tf.function
def test():
  b = tf.math.maximum(0,1)
  if 0 < b <= 2:
   b = b+1
  return b
test().numpy()

OperatorNotAllowedInGraphError: in user code:

    File "<ipython-input-80-cfa5312e70c9>", line 4, in test  *
        if 0 < b <= 2:

    OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Describe the expected behavior
I suppose that it will give 2 as with this little modified version:

@tf.function
def test():
  b = tf.math.maximum(0,1)
  if 0 < b and b <= 2:
   b = b+1
  return b
test().numpy()

2

Contributing

• Do you want to contribute a PR? (yes/no):
  With some hints
• Briefly describe your candidate solution(if contributing):
  I need some pointer

Standalone code to reproduce the issue
Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.

Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

[bhack]

Just to share more debugging info for the failing and working case:

INFO:tensorflow:Source code of <function test at 0x7f06040d5160>:

@tf.function
def test():
  b = tf.math.maximum(0,1)
  if 0 < b and b <= 2:
   b = b+1
  return b


INFO:tensorflow:Transformed <function test at 0x7f06040d5160>:

# coding=utf-8
def tf__test():
    with ag__.FunctionScope('test', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:
        do_return = False
        retval_ = ag__.UndefinedReturnValue()
        b = ag__.converted_call(ag__.ld(tf).math.maximum, (0, 1), None, fscope)

        def get_state():
            return (b,)

        def set_state(vars_):
            nonlocal b
            (b,) = vars_

        def if_body():
            nonlocal b
            b = ag__.ld(b) + 1

        def else_body():
            nonlocal b
            pass
        ag__.if_stmt(ag__.and_(lambda : 0 < ag__.ld(b), lambda : ag__.ld(b) <= 2), if_body, else_body, get_state, set_state, ('b',), 1)
        try:
            do_return = True
            retval_ = ag__.ld(b)
        except:
            do_return = False
            raise
        return fscope.ret(retval_, do_return)

INFO:tensorflow:Source code of <function test at 0x7ff34af56160>:

@tf.function
def test():
  b = tf.math.maximum(0,1)
  if 0 < b <= 2:
   b = b+1
  return b


INFO:tensorflow:Transformed <function test at 0x7ff34af56160>:

# coding=utf-8
def tf__test():
    with ag__.FunctionScope('test', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:
        do_return = False
        retval_ = ag__.UndefinedReturnValue()
        b = ag__.converted_call(ag__.ld(tf).math.maximum, (0, 1), None, fscope)

        def get_state():
            return (b,)

        def set_state(vars_):
            nonlocal b
            (b,) = vars_

        def if_body():
            nonlocal b
            b = ag__.ld(b) + 1

        def else_body():
            nonlocal b
            pass
        ag__.if_stmt(0 < ag__.ld(b) <= 2, if_body, else_body, get_state, set_state, ('b',), 1)
        try:
            do_return = True
            retval_ = ag__.ld(b)
        except:
            do_return = False
            raise
        return fscope.ret(retval_, do_return)

[bhack]

The 2nd autograph is valid so I think that there is a false positive:

from tensorflow.python.autograph.impl import api
ag__ = api._TRANSPILER.get_extra_locals()['ag__']  # pylint:disable=protected-access
def tf__test():
    with ag__.FunctionScope('test', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:
        do_return = False
        retval_ = ag__.UndefinedReturnValue()
        b = ag__.converted_call(ag__.ld(tf).math.maximum, (0, 1), None, fscope)

        def get_state():
            return (b,)

        def set_state(vars_):
            nonlocal b
            (b,) = vars_

        def if_body():
            nonlocal b
            b = ag__.ld(b) + 1

        def else_body():
            nonlocal b
            pass
        ag__.if_stmt(0 < ag__.ld(b) <= 2, if_body, else_body, get_state, set_state, ('b',), 1)
        try:
            do_return = True
            retval_ = ag__.ld(b)
        except:
            do_return = False
            raise
        return fscope.ret(retval_, do_return)

print(tf__test().numpy())

2

[bhack]

The second form seems to pass with:
@tf.function(experimental_autograph_options=tf.autograph.experimental.Feature.EQUALITY_OPERATORS)

As it is converted in the same form as the first impl:

tensorflow/tensorflow/python/autograph/converters/logical_expressions.py

Lines 89 to 90 in 87462bf

	# Repeated comparisons are converted to conjunctions:
	# a < b < c -> a < b and b < c

ag__.if_stmt(ag__.and_(lambda : 0 < ag__.ld(b), lambda : ag__.ld(b) < 2), if_body, else_body, get_state, set_state, ('b',), 1)

/cc @mdanatg

[bhack]

tensorflow/tensorflow/python/autograph/converters/logical_expressions.py

Lines 89 to 90 in 87462bf

	# Repeated comparisons are converted to conjunctions:
	# a < b < c -> a < b and b < c

[mdanatg]

Odd, I thought we had enabled that feature by default. But I think if was left off as Tensor implemented its own __eq__ , __lt__, etc. This use case indicates that we should turn it on anyway.

[bhack]

Odd, I thought we had enabled that feature by default. But I think if was left off as Tensor implemented its own __eq__ , __lt__, etc. This use case indicates that we should turn it on anyway.

Is #55278 (comment) running fine cause we are in eager mode?

[mdanatg]

Yes. When evaluating a < b < c Python calls something in the lines of (a < b).__bool__(), which only works in eager mode.

[bhack]

Yes. When evaluating a < b < c Python calls something in the lines of (a < b).__bool__(), which only works in eager mode.

Yes honestly it is so hard to explain this kind of underline machinery to our users/developers. I needed to point another user to the Cpython source for a "similar" case