def rw_normalize(A):
    """
    Random walk normalization: computes D^⁼1 * A.

    Parameters
    ----------
    A: torch.Tensor
        The matrix to normalize.

    Returns
    -------
    A_norm: torch.FloatTensor
        The normalized adjacency matrix.
    """
    degs = A.sum(dim=1)
    degs[degs == 0] = 1
    return A / degs[:, None]