Tensorflow: TensorArray as Member Variable

Viewed 12

I have an issue with the following code snippet:

import tensorflow as tf    
class Module(tf.Module):
    
    def __init__(self):
        self.cache = tf.TensorArray(size=0, dynamic_size=True, clear_after_read=False, dtype=tf.int32)

    @tf.function
    def write(self):
        for i in range(10):
            x = tf.constant([1,2,3,4,5], dtype=tf.int32)
            self.cache = self.cache.write(self.cache.size(), x)

            tf.print(self.cache.stack())

module = Module()
f = module.write.get_concrete_function()
f()

It gives me the error

OperatorNotAllowedInGraphError: Using a symbolic tf.Tensor as a Python bool is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

However, this works:

import tensorflow as tf
class Module(tf.Module):

    def __init__(self):
        pass

    @tf.function
    def write(self):
        cache = tf.TensorArray(size=0, dynamic_size=True, clear_after_read=False, dtype=tf.int32)
        for i in range(10):
            x = tf.constant([1,2,3,4,5], dtype=tf.int32)
            cache = cache.write(cache.size(), x)

            tf.print(cache.stack())

module = Module()
f = module.write.get_concrete_function()
f()

I presume that TensorArrays are not supported as member variables? Does somebody know how to overcome the issue?

0 Answers
Related