Fixes `random_flax_module` with `flax.linen.BatchNorm`

[juanitorduz]

Fixes https://github.com/pyro-ppl/numpyro/issues/1446

[juanitorduz]

The first two commits are from an old branch 🤦 . We can squash and merge instead.

[juanitorduz]

[juanitorduz]

@fehiepsi should I leave _substitute_default_key in utils.py or is handlers.py a better place?

[fehiepsi]

Leaving it in utils sounds reasonable to me. It is just a workaround for the edge case.

Thanks for fixing the issue!