Missed TensorFlow Dev Summit? Check out the video playlist. Watch recordings

tf.autograph.experimental.set_loop_options

View source on GitHub

Specifies additional arguments to be passed to the enclosing while_loop.

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

The parameters apply to and only to the immediately enclosing loop. It only has effect if the loop is staged as a TF while_loop; otherwise the parameters have no effect.

Usage:

@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 

Also see tf.while_loop.

Args:

  • parallel_iterations: The maximum number of iterations allowed to run in parallel at any given time. Note that this does not guarantee parallel execution.
  • swap_memory: Whether to store intermediate values needed for gradients on the CPU instead of GPU.
  • maximum_iterations: Allows limiting the total number of iterations executed by the loop.
  • shape_invariants: Allows controlling the argument with the same name passed to tf.while_loop. Unlike tf.while_loop, this is a list of (tensor, shape) pairs.