code: def flax_init_rng(rng: jax.Array) -> jax.Array:
      """Returns the rng key that Flax will pass to the Haiku init function."""
      return ParamsRngModule().apply({}, rngs={"params": rng})
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  from haiku._src import base
  from haiku._src import filtering
  from haiku._src import module
  from haiku._src import transform
  from haiku._src.flax import flax_module
  from haiku._src.flax import utils
  from haiku._src.nets import resnet
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: flax_init_rng(rng) and returns: ParamsRngModule