Inconsistent eager/tf.function behavior for rank 0 shape in tf.reshape

[jonas-eschle]

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): Ubuntu 20.04
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.5.0
• Python version: 3.8

Describe the current behavior

The eager version of tf.reshape takes a rank 0 tensor as a shape parameter while the jitted (tf.function decorated) does not.

Describe the expected behavior

It should be consistent, either fail in both or allow in both. To be consistent with other methods, I think it should fail in both.

Contributing

• Do you want to contribute a PR? (yes/no): maybe
• Briefly describe your candidate solution(if contributing): raise an error, such as done in other methods like tf.random.uniform

Standalone code to reproduce the issue
executable example here

import tensorflow as tf

def func():
    return tf.reshape([[42]], 1)


func_jit = tf.function(func=func)

func()  # works
func_jit()  # fails

[kumariko]

I am able to reproduce the error in tensorflow v2.11.0-dev20220829.

Please find the gist for your reference here. Thanks!

[Saduf2019]

@mayou36
Can you please confirm the jit compilation is enabled.

[jonas-eschle]

By "jit" I mean "tf.function" decorated (jitted), not XLA compiled. So it works in eager, but not inside a tf.function.

It seems to happen with and without jit_compiled on, both.

[Saduf2019]

@jonas-eschle
Please confirm if this is still an issue.

[jonas-eschle]

It is. I've updated the gist and it is still broken, even with TF-nightly

[bgdax]

Thank you for raising this issue. We are discussing internally the available options, and we'll update with the result.