def METHOD_NAME(self, grad, var, apply_state=None):
    var_device, var_dtype = var.device, var.dtype.base_dtype
    coefficients = (apply_state or {}).get(
        (var_device, var_dtype)
    ) or self._fallback_apply_state(var_device, var_dtype)
    if self._momentum:
        momentum_var = self.get_slot(var, "momentum")
        return tf.raw_ops.ResourceApplyKerasMomentum(
            var=var.handle,
            accum=momentum_var.handle,
            lr=coefficients["lr_t"],
            grad=grad,
            momentum=coefficients["momentum"],
            use_locking=self._use_locking,
            use_nesterov=self.nesterov,
        )
    else:
        return tf.raw_ops.ResourceApplyGradientDescent(
            var=var.handle,
            alpha=coefficients["lr_t"],
            delta=grad,
            use_locking=self._use_locking,
        )