Add vmap support for rng_bit_generator.

[LenaMartens]

Adds a naive batching rule which loops over the primitive, then stacks the results. This has the advantage that the vmapped result will be the same as calling the primitive n times, but is not very performant.

[froystig]

Side note that something like #7199 might have been useful here.

[froystig]

This is an "inline" loop that will emit N individual primitive calls. Is there a batch dimension length at which jax.lax.map would make sense instead? We can follow up after this PR if there's an answer to that question.

[froystig]

Should we give some kind of warning in the meantime, especially if the batch dimension is large?

[LenaMartens]

Thanks for the comments! I agree that emitting n primitive calls for large n is undesirable, I can follow-up with a jax.lax.map version (which as I understand it is just a scan)

I can maybe do some benchmarking if I get to it, but do you have an intuition as to what batch size is "too large" for the regular for loop? Or should we just always use a scan?

[froystig]

I'm not sure what's "too large" but it might depend on the backend/device. Conceptually we want a map and that alone might be reason enough to switch. One thing we could do is always use an explicit scan for the map and then tune its unroll parameter for performance. This has the additional advantage of not increasing the jaxpr size, only the HLO.