code: def _split_params(
      stacked_params: base.Params,
      num_layers: int,
      name_map: LayerStackTransparencyMapping,
  ) -> base.Params:
      """Splits the stacked parameters."""
  
      def _split(x):
          return [
              jnp.squeeze(s, axis=0)
              for s in jnp.split(x, x.shape[0], axis=0)
          ]
  
      params = {}
      for mod_name, mod_params in stacked_params.items():
          split_mod_params = {
              k: _split(v) for k, v in mod_params.items()
          }
          for i in range(num_layers):
              new_mod_name = name_map.stacked_to_flat(mod_name, i)
              if new_mod_name in params:
                  raise ValueError(
                      f"Found conflicting unstacked module name for {mod_name} at"
                      f" {new_mod_name}."
                  )
              params[new_mod_name] = {
                  k: v[i] for k, v in split_mod_params.items()
              }
  
      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 3 for calling a global function with example usage: _split_params(stacked_params, num_layers, name_map) and returns: params, unknown