def METHOD_NAME(self, var_device, var_dtype, apply_state):
    super().METHOD_NAME(var_device, var_dtype, apply_state)
    apply_state[(var_device, var_dtype)].update(
        dict(
            epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),
            neg_lr_t=-apply_state[(var_device, var_dtype)]["lr_t"],
            zero=tf.zeros((), dtype=tf.int64),
        )
    )