keras/src/optimizers/loss_scale_optimizer.py/LossScaleOptimizer/_stateful_handle_non_finite_grads
class LossScaleOptimizer:    def _stateful_handle_non_finite_grads(self):
        # If any inf or nan in grads, downscale loss and reset counter.
        self.step_counter.assign(0)
        self.dynamic_scale.assign(self.dynamic_scale / 2.0)