<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def compute_covariance(L_aug, Y, k, p):
 """Given label matrix L_aug and labels Y, compute the covariance. Args: L: (np.array {0,1}) [n, d] The augmented (indicator) label matrix k: (int) Cardinality p: (np.array float) [k] The class balance """

    n, d = L_aug.shape
    assert Y.shape[0] == n
    mu = compute_mu(L_aug, Y, k, p)
    return (L_aug.T @ L_aug) / n - mu @ np.diag(p) @ mu.T