code: def set(
      *,
      check_jax_usage: Optional[bool] = None,
      module_auto_repr: Optional[bool] = None,
      restore_flatmap: Optional[bool] = None,
      rng_reserve_size: Optional[int] = None,
  ):
      """Sets the given config option(s).
  
      >>> hk.config.set(module_auto_repr=False)
      >>> hk.Linear(1)
      <...Linear object at ...>
      >>> hk.config.set(module_auto_repr=True)
      >>> hk.Linear(1)
      Linear(output_size=1)
  
      Args:
        check_jax_usage: Checks that jax transforms and control flow are used
          appropriately in Haiku transformed functions.
        module_auto_repr: Can be used to disable the "to string" functionality that
          is part of Haiku's base contructor.
        restore_flatmap: Whether legacy checkpoints should be restored in the old
          FlatMap datatype (as returned by ``to_immtable_dict``), default is to
          restore these as plain dicts.
        rng_reserve_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.
      """
      write(get_config(), **filter_none_values(locals()))
  
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 3 for calling a global function with example usage: set() without return types