tff.learning.templates.reject_non_finite_update

Rejects the update if any non-finite value is in the update.

tff.learning.templates.reject_non_finite_update(
    state: Any, update: Any
) -> tuple[tf.Tensor, _MeasurementsType]

This is the default should_reject_update function used in build_apply_optimizer_finalizer.

state Unused optimzier state.
update The update to be applied to the model's weights by the optimizer.

A tuple of:

  • should_reject (bool tensor): True if the update should be rejected, False otherwise.
  • measurements (OrderedDict): A dict with a single key (update_non_finite) an an integer tensor of whether the update was rejected.