orbit.utils.create_loop_fn

Creates a loop function driven by a Python while loop.

step_fn A function taking a nested structure of tf.data.Iterator or DistributedIterator. There are no constraints on the return value of the function (except that it must be compatible with any reduce_fn provided to the returned loop_fn).

A loop function taking required iterator and num_steps parameters, as well as optional state and reduce_fn parameters for accumulating state over multiple iterations of the loop. See the loop_fn definition below for additional details.