InaccessibleTensorError when running model as a .expand_dims call in the main model code

[ae20cg]

System information
I have written a custom model in TensorFlow as a modification of the transformer model.
M1 Pro 32 GB Macbook Pro 12.1 Monterey(also ran on colab T4 GPU 16 GB)

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• TensorFlow installed from (source or binary): Binary
• TensorFlow version (use command below): 2.7.1 also tried 2.8
• Python version: 3.7.1.2
• CUDA/cuDNN version: N/A
• GPU model and memory: 32 GB M1 PRO, Tesla T4 16 GB

Describe the current behavior
Current when I run my model in eager execution model it is fine, but when I run it with @tf.function wrapper I get an inaccessible tensor error when building the AutoGraph. This is coming in the model call method, when I .expand_dims on the call of another layer class. The error is:

TypeError: <tf.Tensor 'hi_d_transformer/ExpandDims:0' shape=(64, 1, 64) dtype=float32> is out of scope and cannot be used here. Use return values, explicit Python locals or TensorFlow collections to access it.
Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values for more information.

<tf.Tensor 'hi_d_transformer/ExpandDims:0' shape=(64, 1, 64) dtype=float32> was defined here:
    File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
      "__main__", mod_spec)
    File "/usr/lib/python3.7/runpy.py", line 85, in _run_code
      exec(code, run_globals)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py", line 16, in <module>
      app.launch_new_instance()
    File "/usr/local/lib/python3.7/dist-packages/traitlets/config/application.py", line 846, in launch_instance
      app.start()
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelapp.py", line 499, in start
      self.io_loop.start()
    File "/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py", line 132, in start
      self.asyncio_loop.run_forever()
    File "/usr/lib/python3.7/asyncio/base_events.py", line 541, in run_forever
      self._run_once()
    File "/usr/lib/python3.7/asyncio/base_events.py", line 1786, in _run_once
      handle._run()
    File "/usr/lib/python3.7/asyncio/events.py", line 88, in _run
      self._context.run(self._callback, *self._args)
    File "/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py", line 122, in _handle_events
      handler_func(fileobj, events)
    File "/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py", line 300, in null_wrapper
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py", line 452, in _handle_events
      self._handle_recv()
    File "/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py", line 481, in _handle_recv
      self._run_callback(callback, msg)
    File "/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py", line 431, in _run_callback
      callback(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py", line 300, in null_wrapper
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py", line 283, in dispatcher
      return self.dispatch_shell(stream, msg)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py", line 233, in dispatch_shell
      handler(stream, idents, msg)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py", line 399, in execute_request
      user_expressions, allow_stdin)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/ipkernel.py", line 208, in do_execute
      res = shell.run_cell(code, store_history=store_history, silent=silent)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/zmqshell.py", line 537, in run_cell
      return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 2718, in run_cell
      interactivity=interactivity, compiler=compiler, result=result)
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 2822, in run_ast_nodes
      if self.run_code(code, result):
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 2882, in run_code
      exec(code_obj, self.user_global_ns, self.user_ns)
    File "<ipython-input-32-1bb5ee661811>", line 13, in <module>
      train_step(inp, tar)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 915, in __call__
      result = self._call(*args, **kwds)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 963, in _call
      self._initialize(args, kwds, add_initializers_to=initializers)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 786, in _initialize
      *args, **kwds))
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 2983, in _get_concrete_function_internal_garbage_collected
      graph_function, _ = self._maybe_define_function(args, kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 3292, in _maybe_define_function
      graph_function = self._create_graph_function(args, kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 3140, in _create_graph_function
      capture_by_value=self._capture_by_value),
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 1161, in func_graph_from_py_func
      func_outputs = python_func(*func_args, **func_kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 677, in wrapped_fn
      out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 1143, in autograph_handler
      user_requested=True,
    File "<ipython-input-20-949802b44a4a>", line 5, in train_step
      predictions,_ = transformer([inp, tar],
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
      outputs = call_fn(inputs, *args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
      return fn(*args, **kwargs)
    File "<ipython-input-14-ac60aac10b0e>", line 131, in call
      if self.scvs == None:
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1321, in if_stmt
      _py_if_stmt(cond, body, orelse)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1374, in _py_if_stmt
      return body() if cond else orelse()
    File "<ipython-input-14-ac60aac10b0e>", line 133, in call
      scvs = tf.expand_dims(scvs, axis=1)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py", line 1082, in op_dispatch_handler
      return dispatch_target(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/array_ops.py", line 435, in expand_dims_v2
      return gen_array_ops.expand_dims(input, axis, name)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/gen_array_ops.py", line 2354, in expand_dims
      "ExpandDims", input=input, dim=axis, name=name)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/op_def_library.py", line 742, in _apply_op_helper
      attrs=attr_protos, op_def=op_def)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 695, in _create_op_internal
      compute_device)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py", line 3784, in _create_op_internal
      op_def=op_def)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py", line 2175, in __init__
      self._traceback = tf_stack.extract_stack_for_node(self._c_op)

The tensor <tf.Tensor 'hi_d_transformer/ExpandDims:0' shape=(64, 1, 64) dtype=float32> cannot be accessed from here, because it was defined in FuncGraph(name=train_step, id=140104354578128), which is out of scope.

I expect the model to compile, I have tried to move the tf.expand_dims to other layers/subclasses and still get the same bug no matter where. I guess this could be rooted from somewhere else as well but I cannot understand the error.

Standalone code to reproduce the issue

Reproducible test case can be found here: https://colab.research.google.com/gist/ae20cg/fe1577f53f6e68db1e3429643f1e957a/tft2tf2-graph.ipynb

[mohantym]

Hi @ae20cg ! I was able to replicate in 2.7 and 2.8 , But Can you try again after changing empty lists to Tensor arrays ? Attaching relevant threads for reference. R 1 2. Thank you!

[ae20cg]

@mohantym

Since the argument that is causing the code to fail is 'optional', I removed it and the code runs is able to run train. And the tensor that is causing an error does not come from a list.

I have updated the code to get rid of the empty lists, still the same error.

[ae20cg]

I have fixed this issue.

For those this may apply to. This tensor was being called by another layer in the call method. This tensor was then being passed to another layer, however rather than calling it through the call method I was passing it in in initialization.

I changed said method to accept the tensor in the call argument and it fixed the code.

Thanks!

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[kimlemon1]

Hi @ae20cg,

I'm making visual transformer code with TensorFlow 2.x and have the same error like you, so I checked test case you uploaded but I can't find wrong part in the code. Can you tell me where's the problem in the test case code?

I read your comment above about how you solved but I think I didn't understand correctly.

Thank you!