TensorFlow 2.0 - @tf.function 和 tf.unstack TypeError

Question

I have encountered the following problem: by using @tf.function I would like to unstack a tensor along a defined component.

@tf.function
def f1(x):
    y = tf.unstack(x)
    return y 

@tf.function
def f2(x):
    y = tf.unstack(x, axis=0)
    return y

@tf.function
def f3(x):
    y = tf.unstack(x, axis=1)
    return y

x = tf.random.uniform((4,2))
y1 = tf.unstack(x, axis=0) #f2
y2 = tf.unstack(x, axis=1) #f3
y = f1(x) # No problem! (output equal to y1)
z = f2(x) #Problem!
zz = f3(x) #Problem

TypeError: in user code:

<ipython-input-339-c5b8c0b032bb>:8 f2  *
    y = tf.unstack(x, axis=0)

TypeError: 'set' object is not callable

Not sure if it is due to my ignorance with AutoGraph and @tf.function or something else is going wrong. Would appreciate if someone there can make me understand what's going on:-)

Answer 1

I am able to execute your code in Jupyter Notebook both in Tensorflow 1.15.0 and 2.1.0 without any error.

For the benefit of community, below i have mentioned successful run with outputs using TF 2.1.0.

import tensorflow as tf
print(tf.__version__)

@tf.function
def f1(x):
    y = tf.unstack(x)
    return y 

@tf.function
def f2(x):
    y = tf.unstack(x, axis=0)
    return y

@tf.function
def f3(x):
    y = tf.unstack(x, axis=1)
    return y

x = tf.random.uniform((4,2))
y1 = tf.unstack(x, axis=0) #f2
y2 = tf.unstack(x, axis=1) #f3
y = f1(x) # No problem! (output equal to y1)
z = f2(x) 
zz = f3(x)
print(y)
print(z)
print(zz)

Output:

2.1.0
[<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.42976737, 0.00961947], dtype=float32)>, <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.64688444, 0.7597277 ], dtype=float32)>, <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.05788946, 0.5703846 ], dtype=float32)>, <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.5646384 , 0.36961722], dtype=float32)>]
[<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.42976737, 0.00961947], dtype=float32)>, <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.64688444, 0.7597277 ], dtype=float32)>, <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.05788946, 0.5703846 ], dtype=float32)>, <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.5646384 , 0.36961722], dtype=float32)>]
[<tf.Tensor: shape=(4,), dtype=float32, numpy=array([0.42976737, 0.64688444, 0.05788946, 0.5646384 ], dtype=float32)>, <tf.Tensor: shape=(4,), dtype=float32, numpy=array([0.00961947, 0.7597277 , 0.5703846 , 0.36961722], dtype=float32)>]