Support custom prng key

[fehiepsi]

Resolves #1641

[fehiepsi]

cc @junpenglao

[martinjankowiak]

lgtm i take it these are all expected?

  /opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/numpyro/contrib/einstein/steinvi.py:13: DeprecationWarning: jax.random.KeyArray is deprecated. Use jax.Array for annotations, and jax.dtypes.issubdtype(arr, jax.dtypes.prng_key) for runtime detection of typed prng keys.
    from jax.random import KeyArray

test/contrib/test_enum_elbo.py: 60 warnings
  /opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/numpyro/ops/provenance.py:37: DeprecationWarning: jax.linear_util.wrap_init is deprecated. Use jax.extend.linear_util.wrap_init instead.
    wrapped_fun, out_tree = flatten_fun(lu.wrap_init(fn), in_tree)

test/contrib/test_module.py::test_flax_module
  /opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/flax/linen/module.py:77: DeprecationWarning: jax.random.KeyArray is deprecated. Use jax.Array for annotations, and jax.dtypes.issubdtype(arr, jax.dtypes.prng_key) for runtime detection of typed prng keys.
    KeyArray = Union[jax.Array, jax.random.KeyArray]  # pylint: disable=invalid-name

test/contrib/test_module.py::test_flax_module
  /opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/flax/linen/recurrent.py:45: DeprecationWarning: jax.random.KeyArray is deprecated. Use jax.Array for annotations, and jax.dtypes.issubdtype(arr, jax.dtypes.prng_key) for runtime detection of typed prng keys.
    PRNGKey = jax.random.KeyArray

test/contrib/test_module.py::test_flax_module
  /opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/flax/linen/stochastic.py:28: DeprecationWarning: jax.random.KeyArray is deprecated. Use jax.Array for annotations, and jax.dtypes.issubdtype(arr, jax.dtypes.prng_key) for runtime detection of typed prng keys.
    KeyArray = Union[jax.Array, jax.random.KeyArray]

test/contrib/test_module.py::test_haiku_module
  /opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/haiku/_src/dot.py:163: DeprecationWarning: jax.linear_util.transformation is deprecated. Use jax.extend.linear_util.transformation instead.
    @jax.linear_util.transformation```

[fehiepsi]

Thanks, the first two warnings can be fixed. The rest needs to be fixed upstream on flax/haiku.