code: def partition_n(
      fn: Callable[[str, str, T], int],
      structure: Mapping[str, Mapping[str, T]],
      n: int,
  ) -> tuple[Mapping[str, Mapping[str, T]], ...]:
      """Partitions a structure into `n` structures.
  
      For a given set of parameters, you can use :func:`partition_n` to split them
      into ``n`` groups. For example, to split your parameters/gradients by module
      name:
  
      >>> def partition_by_module(structure):
      ...   cnt = itertools.count()
      ...   d = collections.defaultdict(lambda: next(cnt))
      ...   fn = lambda m, n, v: d[m]
      ...   return hk.data_structures.partition_n(fn, structure, len(structure))
  
      >>> structure = {f'layer_{i}': {'w': None, 'b': None} for i in range(3)}
      >>> for substructure in partition_by_module(structure):
      ...   print(substructure)
      {'layer_0': {'b': None, 'w': None}}
      {'layer_1': {'b': None, 'w': None}}
      {'layer_2': {'b': None, 'w': None}}
  
      Args:
        fn: Callable returning which bucket in ``[0, n)`` the given element should
          be output.
        structure: Haiku params or state data structure to be partitioned.
        n: The total number of buckets.
  
      Returns:
        A tuple of size ``n``, where each element will contain the values for which
        the function returned the current index.
      """
      out = [collections.defaultdict(dict) for _ in range(n)]
      for module_name, name, value in traverse(structure):
          i = fn(module_name, name, value)
          assert 0 <= i < n, f"{i} must be in range [0, {n})"
          out[i][module_name][name] = value
      return tuple(data_structures.to_haiku_dict(o) for o in 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 3 for calling a global function with example usage: partition_n(fn, structure, n) and returns: tuple