def dist(module, init_module, p=2, q=2):
    """
    Calculates l_pq distance of the parameter matrix of a layer from the random
    initialization:
        l_p norm of incoming weights to each hidden unit
        l_q norm on the hidden units
    """
    return (
        (module.weight - init_module.weight)
        .view(module.weight.size(0), -1)
        .norm(p=p, dim=1)
        .norm(q)
        .item()
    )

    Calculates l_pq distance of the parameter matrix of a layer from the random
    initialization:
        l_p norm of incoming weights to each hidden unit
        l_q norm on the hidden units
    
grok/metrics.py