import torch


def jitter_node_pos(pos_matrix, scale=1):
    """ 
    Randomly jitter nodes in xyz-direction.
    Args:
        pos_matrix: matrix with xyz-node positions (N x 3).
        scale: scale factor of jittering
    """
    return pos_matrix + (torch.randn(*pos_matrix.shape).numpy() * scale)