tf.function applied to function with non-tensor argument disallows using tensor as bool

[jameshfisher]

Steps to reproduce

1. Install tensorflow_macos
2. Save the following program as dict_arg.py
3. Run with python dict_arg.py

import tensorflow as tf

@tf.function
def dict_arg(arg):
  img = arg['image']
  if tf.random.uniform(()) > 0.5:
    img = img * 2
  return img

print(dict_arg({'image': tf.range(8)}))

Expected behavior

Program runs as if @tf.function had not been applied, printing a tensor.

Actual behavior

Program raises OperatorNotAllowedInGraphError:

$ python dict_arg.py
2021-03-04 10:41:18.793357: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-03-04 10:41:18.794201: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
WARNING:tensorflow:AutoGraph could not transform <function dict_arg at 0x10c9de3a0> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: module 'gast' has no attribute 'Index'
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
Traceback (most recent call last):
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 447, in converted_call
    converted_f = _convert_actual(target_entity, program_ctx)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 284, in _convert_actual
    transformed, module, source_map = _TRANSPILER.transform(entity, program_ctx)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/pyct/transpiler.py", line 286, in transform
    return self.transform_function(obj, user_context)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/pyct/transpiler.py", line 470, in transform_function
    nodes, ctx = super(PyToPy, self).transform_function(fn, user_context)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/pyct/transpiler.py", line 363, in transform_function
    result = self.transform_ast(node, context)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 252, in transform_ast
    node = self.initial_analysis(node, ctx)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 239, in initial_analysis
    node = qual_names.resolve(node)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/pyct/qual_names.py", line 252, in resolve
    return QnResolver().visit(node)
  File "/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ast.py", line 436, in generic_visit
    value = self.visit(value)
  File "/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ast.py", line 445, in generic_visit
    new_node = self.visit(old_value)
  File "/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/pyct/qual_names.py", line 232, in visit_Subscript
    if not isinstance(s, gast.Index):
AttributeError: module 'gast' has no attribute 'Index'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "dict_arg.py", line 10, in <module>
    print(dict_arg({'image': tf.range(8)}))
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 725, in _initialize
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3196, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 966, in wrapper
    return autograph.converted_call(
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 454, in converted_call
    return _fall_back_unconverted(f, args, kwargs, options, e)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 501, in _fall_back_unconverted
    return _call_unconverted(f, args, kwargs, options)
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 478, in _call_unconverted
    return f(*args, **kwargs)
  File "dict_arg.py", line 6, in dict_arg
    if tf.random.uniform(()) > 0.5:
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 885, in __bool__
    self._disallow_bool_casting()
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 488, in _disallow_bool_casting
    self._disallow_when_autograph_enabled(
  File "/Users/jim/tensorflow_macos_venv_2/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 474, in _disallow_when_autograph_enabled
    raise errors.OperatorNotAllowedInGraphError(
tensorflow.python.framework.errors_impl.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.

System details

$ python --version
Python 3.8.2
$ python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"
v1.12.1-44680-gc3fea33a21 2.4.0-rc0
$ sw_vers
ProductName:    macOS
ProductVersion: 11.2.1
BuildVersion:   20D74
$ /Volumes/Macintosh\ HD/usr/sbin/system_profiler SPHardwareDataType
Hardware:

    Hardware Overview:

      Model Name: MacBook Air
      Model Identifier: MacBookAir8,2
      Processor Name: Dual-Core Intel Core i5
      Processor Speed: 1.6 GHz
      Number of Processors: 1
      Total Number of Cores: 2
      L2 Cache (per Core): 256 KB
      L3 Cache: 4 MB
      Hyper-Threading Technology: Enabled
      Memory: 16 GB
...

Notes

• This program behaves as expected with TensorFlow installed via pip on the same machine. It's only the tensorflow_macos distro that has this problem.
• This program behaves as expected in a Colab/hosted notebook environment.
• The program behaves as expected if the img is passed as a direct argument, rather than within a dict.

[zachcoleman]

This issue seems to be from this error: Cause: module 'gast' has no attribute 'Index'. I got around this by downgrading my version of gast==0.3.3.