Fix hk.while_loop carrying state when reserving variable sizes of rng keys.

[copybara-service]

Fix hk.while_loop carrying state when reserving variable sizes of rng keys.

This is already correctly handled in cond/scan/..., but we need to standardize
on the amount of keys in the internal state before running the while_loop and
after running each iteration to make sure our in- and output state have the
same structure.

Also, fix an issue with assert_unchanged_state to take into account
rng_after set might be smaller than rng_before set.