Iterating over tf.Tensor is not allowed: AutoGraph is disabled in this function

Question

I am using tensorflow 2.1 along with python 3.7

The following snippet of code is being used to build a tensorflow graph. The code runs without errors when executed as a standalone python script. (Probably tensorflow is running in eager mode? I am not sure.)

import tensorflow as tf
patches = tf.random.uniform(shape=(1, 10, 50, 300), dtype=tf.dtypes.float32)
s = tf.shape(patches)
patches = [patches[0][x][y] - tf.reduce_mean(patches[0][x][y]) for y in tf.range(s[2]) for x in tf.range(s[1])]

However, the code fails when this is part of a tensorflow graph. I receive the following error: tensorflow.

python.framework.errors_impl.OperatorNotAllowedInGraphError: iterating over tf.Tensor is not allowed: AutoGraph is disabled in this function. Try decorating it directly with @tf.function.

I also added the decorator @tf.function to the method which wraps the above lines of code. It didn't help. I am not sure if I fully understand the meaning of decorating with @tf.function. I also checked that this could be a problem with using python list comprehension inside the tensorflow graph. I am not sure how to use tf.map_fn or tf.while_loop for my case, since I have nested loops.

Thanks in advance!

Answer 1

List comprehensions are not yet supported in autograph. The error that's raised needs to be improved, too. Piling up on https://github.com/tensorflow/tensorflow/issues/32546 should help resolve it sooner.

Until comprehensions are supported, you have to use map_fn, which in this case would look something like this:

def outer_comp(x):
  def inner_comp(y):
    return patches[0][x][y] - tf.reduce_mean(patches[0][x][y])
  return tf.map_fn(inner_comp, tf.range(s[2]), dtype=tf.float32)
patches = tf.map_fn(outer_comp, tf.range(s[1]), dtype=tf.float32)

That said, I believe you can just use reduce_mean directly:

patches = patches - tf.expand_dims(tf.reduce_mean(patches, axis=3), -1)