Splits a PRNG key into num new keys by adding a leading axis.
Parameters:
key (ArrayLike) – a PRNG key (from key, split, fold_in).
num (int | tuple[int, ...]) – optional, a positive integer (or tuple of integers) indicating
the number (or shape) of keys to produce. Defaults to 2.