class_docstr: Wrapper that always initializes weights from a pre-trained checkpoint.
class_docstr_tok: ["Wrapper","that","always","initializes","weights","from","a","pre-trained","checkpoint","."]
code: def init_weights_and_state(self, input_signature):
      super().init_weights_and_state(input_signature)
      if self.init_checkpoint is None:
          return
  
      print("Loading pre-trained weights from", self.init_checkpoint)
      ckpt = tf.train.load_checkpoint(self.init_checkpoint)
  
      def reshape_qkv(name):
          x = ckpt.get_tensor(name)
          return x.reshape((x.shape[0], -1, 64)).swapaxes(0, 1)
  
      def reshape_o(name):
          x = ckpt.get_tensor(name)
          return x.reshape((-1, 64, x.shape[-1]))
  
      def reshape_bias(name):
          x = ckpt.get_tensor(name)
          return x.reshape((-1, 64))
  
      new_w = [
          ckpt.get_tensor("bert/embeddings/word_embeddings"),
          ckpt.get_tensor("bert/embeddings/token_type_embeddings"),
          ckpt.get_tensor("bert/embeddings/position_embeddings")[
              None, ...
          ],
          ckpt.get_tensor("bert/embeddings/LayerNorm/gamma"),
          ckpt.get_tensor("bert/embeddings/LayerNorm/beta"),
      ]
  
      for i in range(12):  # 12 layers
          new_w += [
              reshape_qkv(
                  f"bert/encoder/layer_{i}/attention/self/query/kernel"
              ),
              reshape_qkv(
                  f"bert/encoder/layer_{i}/attention/self/key/kernel"
              ),
              reshape_qkv(
                  f"bert/encoder/layer_{i}/attention/self/value/kernel"
              ),
              reshape_o(
                  f"bert/encoder/layer_{i}/attention/output/dense/kernel"
              ),
              reshape_bias(
                  f"bert/encoder/layer_{i}/attention/self/query/bias"
              ),
              reshape_bias(
                  f"bert/encoder/layer_{i}/attention/self/key/bias"
              ),
              reshape_bias(
                  f"bert/encoder/layer_{i}/attention/self/value/bias"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/attention/output/dense/bias"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/attention/output/LayerNorm/gamma"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/attention/output/LayerNorm/beta"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/intermediate/dense/kernel"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/intermediate/dense/bias"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/output/dense/kernel"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/output/dense/bias"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/output/LayerNorm/gamma"
              ),
              ckpt.get_tensor(
                  f"bert/encoder/layer_{i}/output/LayerNorm/beta"
              ),
          ]
  
      new_w += [
          ckpt.get_tensor("bert/pooler/dense/kernel"),
          ckpt.get_tensor("bert/pooler/dense/bias"),
      ]
  
      for a, b in zip(fastmath.tree_leaves(self.weights), new_w):
          assert (
              a.shape == b.shape
          ), f"Expected shape {a.shape}, got shape {b.shape}"
      self.weights = jax.tree_unflatten(
          jax.tree_structure(self.weights), new_w
      )
      move_to_device = jax.jit(lambda x: x)
      self.weights = jax.tree_map(move_to_device, self.weights)
  
imports: import os
  import pathlib
  import urllib
  import zipfile
  import gin
  import jax
  import tensorflow
  from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 3 for calling the PretrainedBERT obj's underlying member method code with example usage: obj.init_weights_and_state(input_signature) and returns: x, x, x