keras/src/optimizers/loss_scale_optimizer.py/LossScaleOptimizer/_common_apply
class LossScaleOptimizer:    def _common_apply(self, grads, trainable_variables=None):
        finite = self.check_finite(grads)
        ops.cond(
            finite,
            lambda: self._stateful_handle_finite_grads(
                grads, trainable_variables
            ),
            self._stateful_handle_non_finite_grads,
        )