code: def store_haiku_collections(
      scope: flax.core.Scope,
      **collections: HaikuParamsOrState,
  ):
      """Stores the given Haiku params/state in nested collections under scope."""
      for collection_name, collection in collections.items():
          for mod_name, name, value in hkds.traverse(collection):
              subscope = scope
              for part in mod_name.split("/"):
                  subscope = subscope.push(part, reuse=True)
              subscope.put_variable(collection_name, name, value)
  
file_path: dm-haiku/haiku/_src/flax/flax_module.py
imports: from typing import TypeVar, Union
  import flax.core
  import flax.linen
  from haiku._src import filtering
  from haiku._src import transform
  from haiku._src import typing
  from haiku._src.flax import utils
label_desc: Use image node_id 1 for calling a global function with example usage: store_haiku_collections(scope) without return types