Return a stacked array rather than a tuple from next_rng_keys.
Return a stacked array rather than a tuple from next_rng_keys.

This is consistent with JAX's split API.

PiperOrigin-RevId: 383223695

deepmind/dm-haiku,deepmind/dm-haiku