Fix issue with rng check for new-style PRNG keys
For old-style keys, a shape validation is used to confirm that the argument is a single key, not a batch of keys. For new-style keys, this validation has been missing. This change fixes this oversight, and adds some more comprehensive unit tests.
PiperOrigin-RevId: 565123751
