code: def map(  # pylint: disable=redefined-builtin
      fn: Callable[[str, str, InT], OutT],
      structure: Mapping[str, Mapping[str, InT]],
  ) -> Mapping[str, Mapping[str, OutT]]:
      """Maps a function to an input structure accordingly.
  
      >>> params = {'linear': {'w': 1.0, 'b': 2.0}}
      >>> fn = lambda module_name, name, value: 2 * value if name == 'w' else value
      >>> hk.data_structures.map(fn, params)
      {'linear': {'b': 2.0, 'w': 2.0}}
  
      Note: returns a new structure not a view.
  
      Args:
        fn: criterion to be used to map the input data.
          The ``fn`` argument is expected to be a function taking as inputs the
          name of the module, the name of a given entry in the module data bundle
          (e.g. parameter name) and the corresponding data, and returning a new
          value.
        structure: Haiku params or state data structure to be mapped.
  
      Returns:
        All the input parameters or state as mapped by the input fn.
      """
      out = collections.defaultdict(dict)
      for module_name, name, value in traverse(structure):
          out[module_name][name] = fn(module_name, name, value)
      return data_structures.to_haiku_dict(out)
  
file_path: dm-haiku/haiku/_src/filtering.py
imports: import collections
  from collections.abc import Generator, Mapping, MutableMapping
  from typing import Any, Callable, TypeVar
  from haiku._src import data_structures
  from haiku._src import utils
  import jax
label_desc: Use image node_id 5 for calling a global function with example usage: map(fn, structure) and returns: data_structures