tf.autograph.experimental.set_loop_options

>>> @tf.function(autograph=True)
... def f():
...   n = 0
...   for i in tf.range(10):
...     tf.autograph.experimental.set_loop_options(maximum_iterations=3)
...     n += 1
...   return n

@tf.function(autograph=True)
def f():
  v = tf.constant((0,))
  for i in tf.range(3):
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape([None]))]
    )
    v = tf.concat((v, [i]), 0)
  return v