<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_mu(L_aug, Y, k, p):
 """Given label matrix L_aug and labels Y, compute the true mu params. 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

    # Compute mu
    mu = np.zeros((d, k))
    for y in range(1, k + 1):
        L_y = L_aug[Y == y]
        mu[:, y - 1] = L_y.sum(axis=0) / L_y.shape[0]
    return mu