code: def _stack_params(
      split_params: base.Params,
      num_layers: int,
      name_map: LayerStackTransparencyMapping,
  ) -> base.Params:
      """Stacks the split parameters."""
      params = {}
      make_empty_param_stack = lambda: ([None] * num_layers)
  
      for mod_name, mod_params in split_params.items():
          stacked_name_idx = name_map.flat_to_stacked(mod_name)
          if stacked_name_idx is None:
              continue
          stacked_mod_name, idx = stacked_name_idx
          if stacked_mod_name not in params:
              params[stacked_mod_name] = collections.defaultdict(
                  make_empty_param_stack
              )
  
          for k, v in mod_params.items():
              if params[stacked_mod_name][k][idx] is not None:
                  raise ValueError(
                      f"Found conflicting values for param {stacked_mod_name}/{k} at"
                      f" index {idx}."
                  )
              params[stacked_mod_name][k][idx] = v
  
      for mod_name, mod_params in params.items():
          for k, v in mod_params.items():
              if None in v:
                  raise ValueError(
                      f"Couldn't find all params for {mod_name}/{k}: {v}"
                  )
              mod_params[k] = jnp.stack(v, axis=0)
  
      return params
  
imports: import collections
  import functools
  import inspect
  from typing import Any, Callable, Optional, Protocol, Union
  from haiku._src import base
  from haiku._src import lift
  from haiku._src import module
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 4 for calling a global function with example usage: _stack_params(split_params, num_layers, name_map) and returns: params