keras/src/optimizers/loss_scale_optimizer.py/LossScaleOptimizer/_stateful_handle_finite_grads
class LossScaleOptimizer:    def _stateful_handle_finite_grads(self, grads, trainable_variables):
        scale = self.dynamic_scale
        # Unscale gradients.
        unscaled_grads = [
            g if g is None else ops.divide(g, scale) for g in grads
        ]
        self.inner_optimizer.apply(
            unscaled_grads, trainable_variables=trainable_variables
        )

        def upscale():
            self.step_counter.assign(0)
            self.dynamic_scale.assign(self.dynamic_scale * 2.0)

        def increment():
            self.step_counter.assign_add(1)

        # Potentially upscale loss and reset counter.
        ops.cond(
            ops.equal(self.step_counter, self.dynamic_growth_steps - 1),
            upscale,
            increment,
        )