I am trying to generate a matrix that is simply dependent on the resulting output size and some calculation parameters, but I am stuck trying to convert it to a proper tf.function.

So I started rewriting python constructs to tf.range and tf.TensorArray and finally got a result, but everything was zero. A minimal example looks like this:

class Test(tf.Module):
    @tf.function
    def __call__(self):
        ff = tf.TensorArray(tf.int32, size=12)
        for i in tf.range(12):
            ff.write(i,i*2)  # complex calculation here
        return tf.reshape(ff.stack(), shape=(3,4))
t = Test()
t()

output:

<tf.Tensor: shape=(3, 4), dtype=int32, numpy=
array([[0, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0]])>

1. what is happening here? why is the output all zeros?
2. is there a "correct" way to do it? I am trying to avoid vectorizing the whole calculation since I think memory consumption will grow quite a bit, as I will have to store several pre-calculated matrices of the output size in order to correctly calculate the desired result.

CodePudding user response：

You are forgetting to assign the values you write to your TensorArray:

class Test(tf.Module):
    @tf.function
    def __call__(self):
        ff = tf.TensorArray(tf.int32, size=12)
        for i in tf.range(12):
            ff = ff.write(i,i*2)  # complex calculation here
        return tf.reshape(ff.stack(), shape=(3,4))
t = Test()
print(t())

tf.Tensor(
[[ 0  2  4  6]
 [ 8 10 12 14]
 [16 18 20 22]], shape=(3, 4), dtype=int32)