Project: 071_Booster_Tackling_Harmful_Fine-tuning_for_Large_Lan | Function: MMD_loss.forward (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:** Maximum Mean Discrepancy loss computation using Gaussian kernels to measure the distance between two probability distributions, used in RepNoise to match harmful embeddings with Gaussian noise

**KEY TERMS:** {
  "kernels": "Gaussian kernel matrix",
  "XX": "source-source kernel values",
  "YY": "target-target kernel values",
  "XY": "source-target kernel values",
  "YX": "target-source kernel values",
  "guassian_kernel": "Gaussian kernel computation"
}

**CONTEXT CODE:** import torch
from torch import nn

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def register_activation_hook(model):
    activations = {}
    hooks = []
    for name, param in model.named_modules():
        param.name = name
        def _hook(module, __, val):
            activations[module.name] = val
            # print(module.name)
        hooks += [param.register_forward_hook(_hook)]
        
    return activations, hooks 


class MMD_loss(nn.Module):
    def __init__(self, kernel_mul = 2.0, kernel_num = 5):
        super(MMD_loss, self).__init__()
        self.kernel_num = kernel_num
        self.kernel_mul = kernel_mul
        self.fix_sigma = None
        return
    
    def guassian_kernel(self, source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
        n_samples = int(source.size()[0])+int(target.size()[0])
        total = torch.cat([source, target], dim=0)
 
        total0 = total.unsqueeze(0).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
        total1 = total.unsqueeze(1).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
        L2_distance = ((total0-total1)**2).sum(2)
        if fix_sigma:
            bandwidth = fix_sigma
        else:
            bandwidth = torch.sum(L2_distance.data) / (n_samples**2-n_samples)
        bandwidth /= kernel_mul ** (kernel_num // 2)
        bandwidth_list = [bandwidth * (kernel_mul**i) for i in range(kernel_num)]
        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".