title: How to manipulate and return tf.Variable using a for loop over tf.data.Dataset inside function decorated with @tf.function?
body: <p>I am trying to create a function containing a for loop over a TensorFlow Dataset that assigns a new value to a TensorFlow Variable in each iteration. The Variable should also be returned as output of the function. With eager execution enabled, there are no issues, however, in graph mode, some unexpected things seem to happen. Consider the following simple dummy code:</p>
  
  <pre><code>import tensorflow as tf
  
  
  class Test(object):
      def __init__(self):
          self.var = tf.Variable(0, trainable=False, dtype=tf.float32)
          self.increment = tf.constant(1, dtype=tf.float32)
          self.dataset = tf.data.Dataset.from_tensor_slices([0, 1, 2])
  
      @tf.function
      def fn1(self):
          self.var.assign(0)
          for _ in tf.range(3):
              self.var.assign(self.var+self.increment)
              tf.print(self.var)
          tf.print(self.var)
          return self.var
  
      @tf.function
      def fn2(self):
          self.var.assign(0)
          for _ in self.dataset:
              self.var.assign(self.var+self.increment)
              tf.print(self.var)
          tf.print(self.var)
          return self.var
  
      @tf.function
      def fn3(self):
          self.var.assign(0)
          y = self.var
          for _ in self.dataset:
              self.var.assign(self.var+self.increment)
              y = self.var
              tf.print(y)
          tf.print(y)
          return y
  
      @tf.function
      def fn4(self):
          var = 0.0
          for _ in self.dataset:
              var += 1.0
              tf.print(var)
          tf.print(var)
          return var
  
  </code></pre>
  
  <p><code>test.fn1()</code>, <code>test.fn3()</code> and <code>test.fn4()</code> all return the following (desired) output:</p>
  
  <pre><code>1
  2
  3
  3
  &lt;tf.Tensor: shape=(), dtype=float32, numpy=3.0&gt;
  </code></pre>
  
  <p>However, <code>test.fn2()</code> behaves differently:</p>
  
  <pre><code>1
  2
  3
  0
  &lt;tf.Tensor: shape=(), dtype=float32, numpy=0.0&gt;
  </code></pre>
  
  <p>Interestingly, after execution of <code>test.fn2</code>, <code>test.var</code> does seem to contain the correct value:</p>
  
  <pre><code>&lt;tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0&gt;
  </code></pre>
  
  <p>I am not sure why <code>test.fn2</code> fails. Clearly, it is doing some things correctly (as <code>test.var</code> contains the correct value after execution of the function), but it does not deliver the correct result. Can you help me understand what causes this code to fail?</p>
  
  <p>The behavior described above occurs when using TensorFlow 2.1.0 for Python 3.6 on CentOS 7.</p>