hoist named shape logic from threefry's random_bits implementation

[froystig]

The threefry2x32 PRNG's random_bits implementation handles named shapes by folding axis indices into the key. With #6899 in place, we should attempt to hoist this logic out of threefry and apply it uniformly across key implementations. In turn, hopefully RNG implementations need not be concerned about named shapes, only positional ones.

Related:

• implement a wider PRNG #7676, which should benefit by not needing to worry about named shapes.
• Split keys over named batch axis #7708, in which we may want this logic to behave identically to vmap-of-split.