keras/src/optimizers/base_optimizer.py/BaseOptimizer/_clip_by_norm
class BaseOptimizer:    def _clip_by_norm(self, values, axes=None):
        # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
        l2sum = ops.sum(ops.square(values), axes, keepdims=True)
        pred = l2sum > 0
        # Two-tap tf.where trick to bypass NaN gradients
        l2sum_safe = ops.where(pred, l2sum, ops.ones_like(l2sum))
        l2norm = ops.where(pred, ops.sqrt(l2sum_safe), l2sum)
        intermediate = ops.multiply(values, self.clipnorm)
        values_clip = ops.convert_to_tensor(intermediate) / ops.maximum(
            l2norm, self.clipnorm
        )
        return values_clip