Why I cannot using TensorArray.gather() in @tf.function?

Question

Reading from TensorArray:

def __init__(self, size):
    self.obs_buf = tf.TensorArray(tf.float32, size=size, clear_after_read=False)
    self.obs2_buf = tf.TensorArray(tf.float32, size=size, clear_after_read=False)
    self.act_buf = tf.TensorArray(tf.float32, size=size, clear_after_read=False)
    self.rew_buf = tf.TensorArray(tf.float32, size=size, clear_after_read=False)
    self.done_buf = tf.TensorArray(tf.float32, size=size, clear_after_read=False)

def get_sample(self, batch_size):
        idxs = tf.random.uniform(shape=[batch_size], maxval=self.size, dtype=tf.int32)
        tf.print(idxs)
        return self.obs_buf.gather(indices=idxs),     \     # HERE IS THE ISSUE
               self.act_buf.gather(indices=idxs),     \
               self.rew_buf.gather(indices=idxs),     \
               self.obs2_buf.gather(indices=idxs),    \
               self.done_buf.gather(indices=idxs)

Using:

@tf.function
def train(self, rpm, batch_size, gradient_steps):
    for gradient_step in tf.range(1, gradient_steps + 1):
        obs, act, rew, next_obs, done = rpm.get_sample(batch_size)

        with tf.GradientTape() as tape:
        ...

Issue:

Traceback (most recent call last): File ".\main.py", line 130, in rl_training.train() File "C:\Users\user\Documents\Projects\rl-toolkit\rl_training.py", line 129, in train self._rpm, self.batch_size, self.gradient_steps, logging_wandb=self.logging_wandb File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\def_function.py", line 828, in call result = self._call(*args, **kwds) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\def_function.py", line 871, in _call self._initialize(args, kwds, add_initializers_to=initializers) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\def_function.py", line 726, in _initialize *args, **kwds)) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\function.py", line 2969, in _get_concrete_function_internal_garbage_collected graph_function, _ = self._maybe_define_function(args, kwargs) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\function.py", line 3361, in _maybe_define_function graph_function = self._create_graph_function(args, kwargs) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\function.py", line 3206, in _create_graph_function capture_by_value=self._capture_by_value), File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\framework\func_graph.py", line 990, in func_graph_from_py_func func_outputs = python_func(*func_args, **func_kwargs) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\def_function.py", line 634, in wrapped_fn out = weak_wrapped_fn().wrapped(*args, **kwds) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\eager\function.py", line 3887, in bound_method_wrapper return wrapped_fn(*args, **kwargs) File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\framework\func_graph.py", line 977, in wrapper raise e.ag_error_metadata.to_exception(e) tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: in user code:

C:\Users\user\Documents\Projects\rl-toolkit\policy\sac\sac.py:183 update  *
    obs, act, rew, next_obs, done = rpm.get_sample(batch_size)
C:\Users\user\Documents\Projects\rl-toolkit\utils\replay_buffer.py:39 __call__  *
    return self.obs_buf.gather(indices=idxs),                    self.act_buf.gather(indices=idxs),                    self.rew_buf.gather(indices=idxs),                    self.obs2_buf.gather(indices=idxs),                   self.done_buf.gather(indices=idxs)
C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\ops\tensor_array_ops.py:1190 gather  **
    return self._implementation.gather(indices, name=name)
C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\ops\tensor_array_ops.py:861 gather
    return array_ops.stack([self._maybe_zero(i) for i in indices])
C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\framework\ops.py:505 __iter__
    self._disallow_iteration()
C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\framework\ops.py:498 _disallow_iteration
    self._disallow_when_autograph_enabled("iterating over `tf.Tensor`")
C:\Users\user\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\framework\ops.py:476 _disallow_when_autograph_enabled
    " indicate you are trying to use an unsupported feature.".format(task))

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Why I cannot using TensorArray in this context? And what alternatives I have?

does this github.com/tensorflow/tensorflow/issues/31952 help ? — Andrey
Sorry, but no becouse I have trouble with tf.TensorArray.gather() not tf.gather() ..... +0 solution is not working in this case. — Bc. Martin Kubovčík

Bc. Martin Kubovčík Bc. Martin Kubovčík · Accepted Answer · 2021-06-24T11:30:08

0

votes

Solved here. Must be used tf.Variable instead of tf.TensorArray.

Why I cannot using TensorArray.gather() in @tf.function?

1 Answers