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.

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

Returns
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.