title: How do I efficiently track the history of a tensorflow tensor?
body: <p>Is this pattern of using a <code>TensorArray</code> an efficient way of keeping track of a tensor's history?  Is everything inside the inner loop done on the GPU without transmitting anything to the CPU?   How can I verify this?</p>
  
  <pre><code>import tensorflow as tf
  
  
  with tf.device('/device:GPU:0'):
      @tf.function
      def f(x, y):
          return y, x + y
  
      x_array = tf.TensorArray(tf.float32, 0, dynamic_size=True,
                               clear_after_read=False)
      y_array = tf.TensorArray(tf.float32, 0, dynamic_size=True,
                               clear_after_read=False)
  
      x = tf.Variable([1.0])
      y = tf.Variable([1.0])
      x_array.write(0, x)
      y_array.write(0, y)
  
      for i in tf.range(10):
          x = x_array.read(i)
          y = y_array.read(i)
          new_x, new_y = f(x, y)
          x_array.write(i + 1, new_x)
          y_array.write(i + 1, new_y)
  
      print(x_array.stack())
      print(y_array.stack())
  </code></pre>
  
  <p>This is what I'm actually trying to do, but it doesn't even run:</p>
  
  <pre><code>import tensorflow as tf
  
  
  with tf.device('/device:GPU:0'):
      @tf.function
      def f(x, y):
          return y, x + y
  
      @tf.function
      def g(n):
          for i in tf.range(n):
              x = x_array.read(i)
              y = y_array.read(i)
              new_x, new_y = f(x, y)
              x_array.write(i + 1, new_x)
              y_array.write(i + 1, new_y)
  
      x_array = tf.TensorArray(tf.float32, 0, dynamic_size=True,
                               clear_after_read=False)
      y_array = tf.TensorArray(tf.float32, 0, dynamic_size=True,
                               clear_after_read=False)
  
      x = tf.Variable([1.0])
      y = tf.Variable([1.0])
      x_array.write(0, x)
      y_array.write(0, y)
  
      g(tf.constant(10))
  
      print(x_array.stack())
      print(y_array.stack())
  </code></pre>