code: def merge_params_dim_0(
      param: Union[nn.Parameter, torch.Tensor],
      key: str,
      p1: Dict[str, torch.Tensor],
      p2: Dict[str, torch.Tensor],
  ):
      """
      ### Load a parameter by merging the partitions along first dimension
  
      :param param: is the parameter
      :param key: is the name of the parameter
      :param p1: first partition dictionary
      :param p2: second partition dictionary
      """
      w1, w2 = p1[key], p2[key]
      param.data[: w1.shape[0]] = w1
      param.data[w1.shape[0] :] = w2
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/neox/checkpoint.py
imports: from pathlib import Path
  from typing import Dict, Union, Tuple, Optional
  import torch
  from torch import nn
  from labml import monit, lab, logger
  from labml.logger import Text, inspect
  from labml.utils.download import download_file
label_desc: Use image node_id 5 for calling a global function with example usage: merge_params_dim_0(param, key, p1, p2) without return types