<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def key_expand(key, Nk=4):
 """Expand the key into the round key."""

    assert Nk in {4, 6, 8}
    Nr = Nk + 6

    key = key.reshape(Nk, 32)
    rkey = exprzeros(4*(Nr+1), 32)

    for i in range(Nk):
        rkey[i] = key[i]

    for i in range(Nk, 4*(Nr+1)):
        if i % Nk == 0:
            rkey[i] = rkey[i-Nk] ^ subword(rotword(rkey[i-1])) ^ RCON[i//Nk].zext(32-8)
        elif Nk > 6 and i % Nk == 4:
            rkey[i] = rkey[i-Nk] ^ subword(rkey[i-1])
        else:
            rkey[i] = rkey[i-Nk] ^ rkey[i-1]

    return rkey