Project: 042_Robustness_Inspired_Graph_Backdoor_Defense | Function: prediction_variance_calculation (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:** Calculate prediction variance using KL divergence between predictions and original prediction for robustness measurement

**KEY TERMS:** {
  "predictions": "pk i prediction logits",
  "ori_predict": "pi original predictions",
  "deviations": "KL divergence values",
  "summed_deviations": "s(i) prediction variance"
}

**CONTEXT CODE:** device = None
drop_ratio = 0.5

def sample_noise_all(edge_index, edge_weight, device):
    """
    Random edge dropping mechanism for graph backdoor defense.
    Implements the paper's noise injection strategy by randomly dropping edges.
    
    Args:
        edge_index: Edge indices tensor [2, num_edges]
        edge_weight: Edge weights tensor [num_edges]
        device: PyTorch device
        
    Returns:
        noisy_edge_index: Edge indices after dropping
        noisy_edge_weight: Edge weights after dropping
    """
    # Ensure inputs are on the correct device
    edge_index = edge_index.to(device)
    if edge_weight is None:
        edge_weight = torch.ones(edge_index.size(1), device=device)
    else:
        edge_weight = edge_weight.to(device)

    # Generate mask for edge dropping
    drop_mask = Bernoulli(1 - drop_ratio).sample(edge_weight.size()).bool()

    # Apply mask to edges
    noisy_edge_index = edge_index[:, drop_mask]
    noisy_edge_weight = edge_weight[drop_mask]

    # Get node degrees
    node_degrees = torch.zeros(edge_index.max() + 1, device=device)
    node_degrees.index_add_(0, noisy_edge_index[0], torch.ones(noisy_edge_index.size(1), device=device))
    # print('degree', node_degrees)

    # Restore edges for isolated nodes
    isolated_nodes = node_degrees == 0
    # print('isolated_nodes', isolated_nodes)
    if isolated_nodes.any():
        potential_restore_edges = isolated_nodes[edge_index[0]]
        # print('potential_restore_edges', potential_restore_edges)
        restore_edges = edge_index[:, potential_restore_edges]
        noisy_edge_index = torch.cat([noisy_edge_index, restore_edges], dim=1)
        restored_weights = torch.ones(restore_edges.size(1), device=device)
        noisy_edge_weight = torch.cat([noisy_edge_weight, restored_weights], dim=0)

    return noisy_edge_index, noisy_edge_weight


def prediction_variance_calculation(predictions, ori_predict, epsilon=1e-8):
    """
    Calculate prediction variance using KL divergence between predictions and original prediction.
    Corresponds to the paper's robust training objective that measures prediction stability.
    
    Args:
        predictions: List of prediction tensors from multiple noise samples
        ori_predict: Original prediction tensor (without noise)
        epsilon: Small value to avoid log(0)
        
    Returns:
        summed_deviations: Summed KL divergence for each node [num_nodes, num_classes]
    """
    deviations = []
    for sub_pred in predictions:
        sub_pred += epsilon
        deviation = F.kl_div(sub_pred.log(), ori_predict, reduce=False)
        deviations.append(deviation)

    summed_deviations = torch.zeros_like(deviations[0]).to(deviations[0].device)
    for deviation in deviations:
        ##### summed deviations for each node #####
        summed_deviations += deviation
    
    return summed_deviations


if __name__ == "__main__":
    # args = parser.parse_args()
    args = parser.parse_known_args()[0]
    # args.cuda =  not args.no_cuda and torch.cuda.is_available()
    device = torch.device(('cuda:{}' if torch.cuda.is_available() else 'cpu').format(args.device_id))

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    print(args)

    #### dataset ####

    from torch_geometric.utils import to_undirected
    import torch_geometric.transforms as T
    transform = T.Compose([T.NormalizeFeatures()])

    if(args.dataset == 'Cora' or args.dataset == 'Pubmed'):
        dataset = Planetoid(root='./data/', \
                            name=args.dataset,\
                            transform=transform)
    elif(args.dataset == 'ogbn-arxiv'):
        from ogb.nodeproppred import PygNodePropPredDataset
        # Download and process data at './dataset/ogbg_molhiv/'
        dataset = PygNodePropPredDataset(name = 'ogbn-arxiv', root='./data/')
        split_idx = dataset.get_idx_split()

    data = dataset[0].to(device)

    if(args.dataset == 'ogbn-arxiv'):
        nNode = data.x.shape[0]
        setattr(data,'train_mask',torch.zeros(nNode, dtype=torch.bool).to(device))
        # dataset[0].train_mask = torch.zeros(nEdge, dtype=torch.bool).to(device)
        data.val_mask = torch.zeros(nNode, dtype=torch.bool).to(device)
        data.test_mask = torch.zeros(nNode, dtype=torch.bool).to(device)
        data.y = data.y.squeeze(1)



    #### preprocess data ####

    from utils import get_split
    data, idx_train, idx_val, idx_clean_test, idx_atk = get_split(args,data,device)

    from torch_geometric.utils import to_undirected
    from utils import subgraph
    data.edge_index = to_undirected(data.edge_index)
    train_edge_index,_, edge_mask = subgraph(torch.bitwise_not(data.test_mask),data.edge_index,relabel_nodes=False)
    mask_edge_index = data.edge_index[:,torch.bitwise_not(edge_mask)]

    # from sklearn_extra import cluster
    from models.backdoor import Backdoor
    from models.construct import model_construct

    unlabeled_idx = (torch.bitwise_not(data.test_mask)&torch.bitwise_not(data.train_mask)).nonzero().flatten()

    import ast
    with open(args.idx_attach, 'r') as file:
        content = file.read()
        idx_attach = ast.literal_eval(content)

    idx_attach = torch.tensor(idx_attach)
    idx_attach = idx_attach.to(device)

    #### load trigger generator and poisoned data ####
    trigger_generator_address = args.trigger_generator_address
    poison_x = torch.load(args.poison_x).to(device)
    poison_edge_index = torch.load(args.poison_edge_index).to(device)
    poison_edge_weights = torch.load(args.poison_edge_weights).to(device)
    poison_labels = torch.load(args.poison_labels).to(device)

    #### trigger generator ####
    model = Backdoor(args,device)


Complete the function implementation replacing "raise NotImplementedError".