Random rounding float32 to bfloat16?

[luyug]

The gopher paper mentioned that they found it best to do "random rounding" of parameters from fp32 to bf16 in the forward pass. (https://arxiv.org/pdf/2112.11446.pdf; page 53) I wonder what is the best way to do this in jax?

[jakevdp]

Looks like a duplicate of #20551