import torch


def get_distorted_x(x, prob, corrupt=True):
    """
        Distorts the node feature matrix (X) either through corruption or through blanking out
        - Corruption: Add uniform random noise in [0,1] to the node feature matrix
        - Blanking out: Replace entries in the node feature matrix with zeros

            Parameters:
                x : torch.tensor of shape (num_nodes, num_node_features)
                    Node feature matrix
                prob : float
                    What fraction of entries in the node feature matrix must be distorted
                corrupt : bool, optional
                    Whether to distort X through corruption or blanking out
                        Default value: True (corruption)
        
            Returns:
                Distorted node feature matrix : torch.tensor of shape (num_nodes, num_node_features)
    """
    assert 0 <= prob <= 1, ValueError(f"Expected input in range [0,1], got {prob} instead")
    U = torch.rand_like(x)
    x_n = x + U if corrupt else torch.zeros_like(x)
    return torch.where(U < prob, x_n, x)