code: def _create_partitions(
      module_sparsity: ModuleSparsity, params: hk.Params
  ) -> tuple[Sequence[hk.Params], Sequence[jax.Array], hk.Params]:
      """Partition params based on sparsity_predicate_map.
  
      Args:
        module_sparsity:  A Sequence of (Predicate, float) pairs.  Predicate
          functions take module_name, name, value as arguments. The floats are the
          sparsity level to apply to leaves matching Predicate.
        params: A Haiku param.
  
      Returns:
        A tuple containing:
          - A list of len(module_sparsity), where each element is a disjoint subset
            of the `params` to be pruned.
          - A list of len(module_sparsity) where each element is the sparsity level.
          - The remaining elements of `params` not being pruned such that the union
            of the first list and this element contains the elements of `params`.
      """
      list_of_trees = []
      sparsity_list = []
  
      tail = params
      # Greedily match so that no parameter can be matched more than once
      for predicate, sparsity in module_sparsity:
          head, tail = hk.data_structures.partition(predicate, tail)
          list_of_trees.append(head)
          sparsity_list.append(sparsity)
  
      return list_of_trees, sparsity_list, tail
  
imports: from collections.abc import Iterator, Mapping, Sequence
  import functools
  from typing import Callable
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: _create_partitions(module_sparsity, params) and returns: list_of_trees,  sparsity_list,  tail