efficient untrue batching of random_bit_generator

[froystig]

The batching rule for the random_bit_generator primitive, over a batch of keys, emits a loop (via lax.map):

jax/jax/_src/lax/control_flow/loops.py

Lines 2012 to 2024 in 42ae843

	def _rng_bit_generator_batching_rule(batched_args, batch_dims, *, shape, dtype, algorithm):
	"""Calls RBG in a loop and stacks the results."""
	key, = batched_args
	bd, = batch_dims
	if bd is batching.not_mapped:
	return lax.rng_bit_generator_p.bind(key, shape=shape, dtype=dtype,
	algorithm=algorithm), (None, None)
	key = batching.moveaxis(key, bd, 0)
	map_body = lambda k: lax.rng_bit_generator_p.bind(k, shape=shape, dtype=dtype, algorithm=algorithm)
	stacked_keys, stacked_bits = map(map_body, key)
	return (stacked_keys, stacked_bits), (0, 0)
	batching.primitive_batchers[lax.rng_bit_generator_p] = _rng_bit_generator_batching_rule # type: ignore

This is a workaround to the corresponding RandomBitGenerator HLO not being batchable. But looping violates the operational expectations of vmap, namely that everything is vectorized. And downstream, the surprise performance hit when switching RNG implementations isn't great.

We could consider a few options:

1. Emit an unrolled loop. Drawbacks: grows the program size with the batch size.
2. Generate a batch of random numbers from a single key in the batch, dropping the remaining keys in the batch. Drawbacks: this violates vmap semantics considering the random values generated, although the output is "statistically" the same in a sense.

Let's try number 2.

The RBG operation is already non-portable across platforms and XLA flags. In some cases the random generation is affected by sharding. So arguably, callers opting into RBG RNGs already expect unusual semantics. By contrast, it's uncommon that anyone expects the performance hit.

cc @mattjj, @dlwh

[froystig]

This ought to address #16792