Project: 042_Robustness_Inspired_Graph_Backdoor_Defense | Function: sample_noise_all (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:** Random edge dropping mechanism for graph backdoor defense that implements noise injection strategy

**KEY TERMS:** {
  "drop_mask": "ϵij edge drop noise",
  "drop_ratio": "β probability",
  "noisy_edge_index": "A ⊕ ϵ perturbed graph",
  "isolated_nodes": "nodes with degree 0"
}

**CONTEXT CODE:**                     help="number of poisoning nodes relative to the full graph")
# defense setting
parser.add_argument('--defense_mode', type=str, default="none",
                    choices=['prune', 'isolate', 'none','reconstruct'],
                    help="Mode of defense")
parser.add_argument('--prune_thr', type=float, default=0.8,
                    help="Threshold of prunning edges")
parser.add_argument('--target_loss_weight', type=float, default=1,
                    help="Weight of optimize outter trigger generator")
parser.add_argument('--weight_target', type=float, default=1,
                    help="Weight of optimize outter trigger generator")
parser.add_argument('--weight_ood', type=float, default=1,
                    help="Weight of optimize outter trigger generator")
parser.add_argument('--weight_targetclass', type=float, default=1,
                    help="Weight of optimize outter trigger generator")
parser.add_argument('--outter_size', type=int, default=4096,
                    help="Weight of optimize outter trigger generator")
parser.add_argument('--homo_loss_weight', type=float, default=100,
                    help="Weight of optimize similarity loss")
parser.add_argument('--homo_boost_thrd', type=float, default=0.8,
                    help="Threshold of increase similarity")
# attack setting
parser.add_argument('--dis_weight', type=float, default=1,
                    help="Weight of cluster distance")
parser.add_argument('--selection_method', type=str, default='none',
                    choices=['loss','conf','cluster','none','cluster_degree'],
                    help='Method to select idx_attach for training trojan model (none means randomly select)')
parser.add_argument('--test_model', type=str, default='GCN',
                    choices=['GCN','GAT','GraphSage','GIN'],
                    help='Model used to attack')
parser.add_argument('--evaluate_mode', type=str, default='1by1',
                    choices=['overall','1by1'],
                    help='Model used to attack')

parser.add_argument('--trigger_generator_address', type=str, default='./model_weights_cora.pth')
parser.add_argument('--poison_x', type=str, default='poison_x_cora.pt')
parser.add_argument('--poison_edge_index', type=str, default='poison_edge_index_cora.pt')
parser.add_argument('--poison_edge_weights', type=str, default='poison_edge_weights_cora.pt')
parser.add_argument('--poison_labels', type=str, default='poison_labels_cora.pt')
parser.add_argument('--idx_attach', type=str, default='Cora_UGBA.txt')

# GPU setting
parser.add_argument('--device_id', type=int, default=3,
                    help="Threshold of prunning edges")

# Global variables for testing
args = None
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()


Complete the function implementation replacing "raise NotImplementedError".