Advanced options for orbit.StandardTrainer.

use_tf_function A boolean indicating whether to apply tf.function to the training loop. This will only affect the body of the loop (involving train_step); train_loop_begin and train_loop_end will always be run in eager mode.
use_tf_while_loop A boolean indicating whether to run the training loop using a tf.while_loop. If True, use_tf_function must also be True.
use_tpu_summary_optimization A boolean indicating whether to enable a performance optimization for summaries in TPUs. Writing summaries conditionally with outside compilation on TPUs can be extremely slow. If True, this optimization creates two tf.functions with two XLA programs (one with summary calls, and one without). The program with summaries runs only for one step when summaries should be recorded.



use_tf_function True
use_tf_while_loop True
use_tpu_summary_optimization False