tf.autograph.experimental.set_loop_options

Указывает дополнительные аргументы, которые необходимо передать в окружающий цикл while_loop.

tf.autograph.experimental.set_loop_options(
    parallel_iterations=UNSPECIFIED, swap_memory=UNSPECIFIED,
    maximum_iterations=UNSPECIFIED, shape_invariants=UNSPECIFIED
)

Параметры применяются только к ближайшему окружающему циклу. Они имеют эффект только если цикл размещён как TF while_loop; в противном случае параметры не имеют эффекта.

Использование:

@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

См. также tf.while_loop.