class_docstr: ## Transformer Layer
class_docstr_tok: ["#","#","Transformer","Layer"]
docstr: Code to load the checkpoint
code: def load_state(
      self, p1: Dict[str, torch.Tensor], p2: Dict[str, torch.Tensor]
  ):
      """
      Code to load the checkpoint
      """
      with monit.section("Load transformer layer"):
          # Attention output transform
          checkpoint.merge_params_sum(
              self.attention.output.bias, "attention.dense.bias", p1, p2
          )
          checkpoint.merge_params_dim_1(
              self.attention.output.weight,
              "attention.dense.weight",
              p1,
              p2,
          )
  
          # Attention query, key and value transform
          checkpoint.merge_params_dim_0(
              self.attention.qkv_lin.bias,
              "attention.query_key_value.bias",
              p1,
              p2,
          )
          checkpoint.merge_params_dim_0(
              self.attention.qkv_lin.weight,
              "attention.query_key_value.weight",
              p1,
              p2,
          )
  
          # Layer norm before attention
          checkpoint.merge_params_duplicate(
              self.pre_ln_attn.bias, "input_layernorm.bias", p1, p2
          )
          checkpoint.merge_params_duplicate(
              self.pre_ln_attn.weight, "input_layernorm.weight", p1, p2
          )
  
          # FFN second transform
          checkpoint.merge_params_dim_0(
              self.ffn.dense_h_h4.bias, "mlp.dense_h_to_4h.bias", p1, p2
          )
          checkpoint.merge_params_dim_0(
              self.ffn.dense_h_h4.weight,
              "mlp.dense_h_to_4h.weight",
              p1,
              p2,
          )
  
          # FFN first transform
          checkpoint.merge_params_sum(
              self.ffn.dense_h4_h.bias, "mlp.dense_4h_to_h.bias", p1, p2
          )
          checkpoint.merge_params_dim_1(
              self.ffn.dense_h4_h.weight,
              "mlp.dense_4h_to_h.weight",
              p1,
              p2,
          )
  
          # Layer norm before FFN
          checkpoint.merge_params_duplicate(
              self.pre_ln_ffn.bias,
              "post_attention_layernorm.bias",
              p1,
              p2,
          )
          checkpoint.merge_params_duplicate(
              self.pre_ln_ffn.weight,
              "post_attention_layernorm.weight",
              p1,
              p2,
          )
  
imports: import copy
  import math
  from typing import Dict, Optional, Set, Callable, Any, Generator, Tuple
  import torch
  from torch import nn
  from torch.cuda.amp import autocast
  from labml import monit, logger
  from labml.logger import Text
  from labml_nn.neox import checkpoint
  from labml_nn.neox.utils.cache import get_cache
label_desc: Use image node_id 3 for calling the TransformerLayer obj's underlying member method code with example usage: obj.load_state(p1, p2) without return types