code: def params_or_state_dict(
      module_name: str,
      submodules: set[str],
      which: str,
  ) -> Mapping[str, jnp.ndarray]:
      """Returns module parameters or state for the given module or submodules."""
      assert which in ("params", "state")
      out = {}
      frame = base.current_frame()
      for their_module_name, bundle in getattr(frame, which).items():
          if (
              their_module_name == module_name
              or their_module_name.startswith(module_name + "/")
              or their_module_name in submodules
          ):
              for name, value in bundle.items():
                  fq_name = their_module_name + "/" + name
                  out[fq_name] = (
                      value.current if which == "state" else value
                  )
      return out
  
imports: from collections.abc import Mapping
  import contextlib
  import functools
  import inspect
  import re
  from typing import Any, Callable, ContextManager, NamedTuple, Optional, Protocol, TypeVar
  from haiku._src import base
  from haiku._src import config
  from haiku._src import data_structures
  from haiku._src import utils
  import jax
  import jax.numpy
label_desc: Use image node_id 13 for calling a global function with example usage: params_or_state_dict(module_name, submodules, which) and returns: out