Project: 053_On_the_Role_of_Attention_Heads_in_Large_Language_M | Function: kl_divergence (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:** Computes Kullback-Leibler divergence between probability distributions to measure the impact of attention head ablation on model safety, serving as the core mathematical computation for the SHIPS metric

**KEY TERMS:** {
  "base_pd": "p(qH; θO)",
  "mask_pd": "p(qH; θO \\ θh^l_i)",
  "torch.sum(base_pd * torch.log(base_pd / mask_pd))": "D_KL computation"
}

**CONTEXT CODE:** import torch


def kl_divergence(base_pd, mask_pd):
    raise NotImplementedError("Function implementation has been masked for testing")

Complete the function implementation replacing "raise NotImplementedError".