keras/src/optimizers/loss_scale_optimizer.py/LossScaleOptimizer/check_finite
class LossScaleOptimizer:    def check_finite(self, grads):
        tensor_grads = [g for g in grads if g is not None]
        finite_grads = [ops.all(ops.isfinite(g)) for g in tensor_grads]
        return ops.all(ops.convert_to_tensor(finite_grads))