Project: 071_Booster_Tackling_Harmful_Fine-tuning_for_Large_Lan | Function: rep_noise_loss (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implementation of RepNoise baseline method that perturbs harmful data embeddings using Maximum Mean Discrepancy (MMD) loss to degrade harmful information by reducing embedding distribution to normalized Gaussian noise

**KEY TERMS:** {
  "mmd_loss": "MMD regularizer",
  "noise_loss": "MMD loss between harmful embeddings and Gaussian noise",
  "harmful_activations": "harmful data embeddings",
  "gaussian": "normalized Gaussian noise",
  "beta": "noise regularization weight",
  "alpha": "harmful loss weighting"
}

**CONTEXT CODE:**         kernel_val = [torch.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
        return sum(kernel_val)
 
    def forward(self, source, target, xy_only=False):
        raise NotImplementedError("Function implementation has been masked for testing")
    
def masked_token_ce_loss(
    logits,
    labels,
    mask
):
    # # Shift so that tokens < n predict n
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    # # apply mask
    # # shifted masks
    shift_logit_mask = mask[..., :-1].contiguous()
    expanded_mask = shift_logit_mask.unsqueeze(-1).expand(-1, -1, shift_logits.size(-1))
    shift_label_mask = mask[..., 1:].contiguous()
    shift_logits = shift_logits * expanded_mask
    shift_labels = shift_labels * shift_label_mask
    shift_labels[shift_labels == 0] = -100
    shift_labels = shift_labels.type(torch.LongTensor)
    loss_fct = torch.nn.CrossEntropyLoss()
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)).to(DEVICE), shift_labels.view(-1).to(DEVICE))
    # loss_fct = torch.nn.CrossEntropyLoss()
    # loss = loss_fct(logits.view(-1, logits.size(-1)).to(DEVICE), labels.view(-1).to(DEVICE))
    return loss


def adapt_dimension_b2a(a,b):
    import torch.nn.functional as F
    import copy 
    b= copy.deepcopy(b)
    a= copy.deepcopy(a)
    # Handle the first dimension
    if b.shape[0] < a.shape[0]:
        padding_first_dim = (0, 0, 0, a.shape[0] - b.shape[0])
        b = F.pad(b, padding_first_dim, mode='constant', value=0)
    elif b.shape[0] > a.shape[0]:
        b = b[:a.shape[0], :]

    # Handle the second dimension
    if b.shape[1] < a.shape[1]:
        padding_second_dim = (0, a.shape[1] - b.shape[1])
        b = F.pad(b, padding_second_dim, mode='constant', value=0)
    elif b.shape[1] > a.shape[1]:
        b = b[:, :a.shape[1]]
    return b

def rep_noise_loss(
    model, harmful_batch, harmless_batch=None,
    beta=0.001, alpha=1
):
    """ Calculate the representation noise loss

    Args:
        model (Pytorch Model): The model to calculate the loss, i.e. an LLM loaded with Huggingface Transformers
        harmful_batch (Dataloader): the paired harmful batch
        harmless_batch (Dataloader): the paired harmless batch
        beta (float, optional): _description_. Defaults to 0.001.
        alpha (int, optional): _description_. Defaults to 1.
    """
    raise NotImplementedError("Function implementation has been masked for testing")


Complete the function implementation replacing "raise NotImplementedError".