code: def optimize_rng_use(fun):
      """Optimizes a RNG key splitting in ``fun``.
  
      Our strategy here is to use abstract interpretation to run your function
      twice, the first time we use :func:`jax.eval_shape` to avoid spending any
      flops and simply observe how many times you call :func:`~haiku.next_rng_key`.
      We then run your function again, but this time we reserve enough RNG keys
      ahead of time such that we only need to call :func:`jax.random.split` once.
  
      In the following example, we need three random samples for our weight
      matrices in our 3-layer MLP. To draw these samples we use
      :func:`~haiku.next_rng_key` which will split a new key for each sample. By
      using :func:`optimize_rng_use` Haiku will pre-allocate exactly enough RNGs for
      ``f`` to be evaluated by splitting the input key once and only once. For large
      models (unlike this example) this can lead to a reduction in compilation time
      of both ``init`` and ``apply``, with ``init`` seeing a larger expected speedup
      as it performs more RNG key splitting in general.
  
      >>> def f(x):
      ...   net = hk.nets.MLP([300, 100, 10])
      ...   return net(x)
      >>> f = hk.experimental.optimize_rng_use(f)
      >>> f = hk.transform(f)
      >>> params = f.init(jax.random.PRNGKey(42), jnp.ones([1, 1]))
  
      Args:
        fun: A function to wrap.
  
      Returns:
        A function that applies ``fun`` but only requires one call to
        :func:`jax.random.split` by Haiku.
      """
  
      @functools.wraps(fun)
      def wrapper(*args, **kwargs):
          base.assert_context("optimize_rng_use")
  
          # Extract all current state.
          frame = base.current_frame()
          params = frame.params or None
          if params is not None:
              params = data_structures.to_haiku_dict(params)
          state = frame.state or None
          if state is not None:
              state = base.extract_state(state, initial=True)
          rng = frame.rng_stack.peek()
          if rng is not None:
              rng = rng.internal_state
  
          def pure_fun(params, state, rng, *args, **kwargs):
              with base.new_context(
                  params=params, state=state, rng=rng
              ):
                  return fun(*args, **kwargs)
  
          with count_hk_rngs_requested() as rng_count_f:
              jax.eval_shape(
                  pure_fun, params, state, rng, *args, **kwargs
              )
          rng_count = rng_count_f()
  
          if rng_count:
              base.current_frame().rng_stack.peek().reserve(rng_count)
          return fun(*args, **kwargs)
  
      return wrapper
  
imports: import contextlib
  import functools
  from haiku._src import base
  from haiku._src import data_structures
  import jax
label_desc: Use image node_id 2 for calling a global function with example usage: optimize_rng_use(fun) and returns: wrapper, fun, fun