I recently came across this nice one-liner alternative:

import itertools

key_seq = map(jax.random.key, itertools.count())
key = next(key_seq)

which uses the infinite iterator count.