code: def partition(
      predicate: Callable[[str, str, jax.Array], bool],
      structure: Mapping[str, Mapping[str, T]],
  ) -> tuple[
      Mapping[str, Mapping[str, T]], Mapping[str, Mapping[str, T]]
  ]:
      """Partitions the input structure in two according to a given predicate.
  
      For a given set of parameters, you can use :func:`partition` to split them:
  
      >>> params = {'linear': {'w': None, 'b': None}}
      >>> predicate = lambda module_name, name, value: name == 'w'
      >>> weights, biases = hk.data_structures.partition(predicate, params)
      >>> weights
      {'linear': {'w': None}}
      >>> biases
      {'linear': {'b': None}}
  
      Note: returns new structures not a view.
  
      Args:
        predicate: criterion to be used to partition the input data.
          The ``predicate`` argument is expected to be a boolean 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.
        structure: Haiku params or state data structure to be partitioned.
  
      Returns:
        A tuple containing all the params or state as partitioned by the input
          predicate. Entries matching the predicate will be in the first structure,
          and the rest will be in the second.
      """
      f = lambda m, n, v: int(not predicate(m, n, v))
      return partition_n(f, structure, 2)
  
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 2 for calling a global function with example usage: partition(predicate, structure) and returns: partition_n