ValueError: as_list() is not defined on an unknown TensorShape. during training

[guillaumelorre28]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.17.0-dev20240517

Custom code

Yes

OS platform and distribution

macOS Sonoma

Mobile device

No response

Python version

3.10.13

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

I get the following error when training a transformer for translation:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[9], [line 6](vscode-notebook-cell:?execution_count=9&line=6)
      [1](vscode-notebook-cell:?execution_count=9&line=1) model.compile(optimizer='adam', 
      [2](vscode-notebook-cell:?execution_count=9&line=2)     loss='sparse_categorical_crossentropy', 
      [3](vscode-notebook-cell:?execution_count=9&line=3)     metrics=['accuracy'],
      [4](vscode-notebook-cell:?execution_count=9&line=4)     run_eagerly=False
      [5](vscode-notebook-cell:?execution_count=9&line=5)     )
----> [6](vscode-notebook-cell:?execution_count=9&line=6) model.fit(dataset_train, epochs=10)

File ~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:122, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    [119](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:119)     filtered_tb = _process_traceback_frames(e.__traceback__)
    [120](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:120)     # To get the full stack trace, call:
    [121](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:121)     # `keras.config.disable_traceback_filtering()`
--> [122](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:122)     raise e.with_traceback(filtered_tb) from None
    [123](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:123) finally:
    [124](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:124)     del filtered_tb

File ~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:594, in tree_map(func, tree, is_leaf, none_is_leaf, namespace, *rests)
    [592](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:592) leaves, treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace)
    [593](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:593) flat_args = [leaves] + [treespec.flatten_up_to(r) for r in rests]
--> [594](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:594) return treespec.unflatten(map(func, *flat_args))

ValueError: as_list() is not defined on an unknown TensorShape.

Standalone code to reproduce the issue

Here is a colab link to the code:
https://colab.research.google.com/drive/1BQ4lhaZPP5XGb_IUe-rVacMiFPWCfm9Y?usp=sharing

Running eagerly works fine but an error occurs in graph mode.

Relevant log output

No response

[SuryanarayanaY]

Hi @guillaumelorre28 ,

I have reproduced the reported issue.Attached gist for same.