code: def rng_reserve_size(size: int) -> int:
      """Change amount of RNG keys reserved when calling ``next_rng_key``.
  
      Args:
        size: amount of keys to reserve when splitting off a key
          through ``next_rng_key()``, defaults to 1. Reserving larger blocks of keys
          can improve compilation and run-time of your model. Changing the
          reservation size will change RNG keys returned by ``next_rng_key``, and
          will change the generated random numbers.
      Returns:
        The previous value of the rng_reserve_size setting.
      """
      if size <= 0:
          raise ValueError(
              f"RNG reserve size needs to be more than 0, got {size}."
          )
      config = get_config()
      before, config.rng_reserve_size = config.rng_reserve_size, size
      return before
  
file_path: dm-haiku/haiku/_src/config.py
imports: import contextlib
  import dataclasses
  import threading
  from typing import Optional
label_desc: Use image node_id 9 for calling a global function with example usage: rng_reserve_size(size) and returns: before