orbit.utils.create_tf_while_loop_fn

Creates a loop function compatible with TF's AutoGraph loop conversion.

step_fn A function taking a nested structure of tf.data.Iterator or DistributedIterator. Currently, any return values are ignored.

A loop function taking required iterator and num_steps parameters. If called inside a tf.function, the loop will be converted by AutoGraph into a tf.while_loop construct. See the loop_fn definition below for additional details.