code: def BERT(
      d_model=768,
      vocab_size=30522,
      max_len=512,
      type_vocab_size=2,
      n_heads=12,
      d_ff=3072,
      n_layers=12,
      head=None,
      init_checkpoint=None,
      mode="eval",
  ):
      """BERT (default hparams are for bert-base-uncased)."""
      # TODO(piotrekp1): loading config from model_name
  
      layer_norm_eps = 1e-12
      d_head = d_model // n_heads
  
      word_embeddings = tl.Embedding(vocab_size, d_model)
      type_embeddings = tl.Embedding(type_vocab_size, d_model)
      position_embeddings = tl.PositionalEncoding(max_len, mode=mode)
      embeddings = [
          tl.Select([0, 1, 0], n_in=3),  # Drops 'idx' input.
          tl.Parallel(
              word_embeddings,
              type_embeddings,
              [
                  tl.PaddingMask(),
                  tl.Fn(
                      "Squeeze",
                      lambda x: np.squeeze(x, (1, 2)),
                      n_out=1,
                  ),
              ],
          ),
          tl.Add(),
          position_embeddings,
          tl.LayerNorm(epsilon=layer_norm_eps),
      ]
  
      encoder = []
      for _ in range(n_layers):
          attn = tl.SelfAttention(
              n_heads=n_heads,
              d_qk=d_head,
              d_v=d_head,
              bias=True,
              masked=True,
              mode=mode,
          )
          feed_forward = [tl.Dense(d_ff), tl.Gelu(), tl.Dense(d_model)]
          encoder += [
              tl.Select([0, 1, 1]),  # Save a copy of the mask
              tl.Residual(
                  attn, AddBias()
              ),  # pylint: disable=no-value-for-parameter
              tl.LayerNorm(epsilon=layer_norm_eps),
              tl.Residual(*feed_forward),
              tl.LayerNorm(epsilon=layer_norm_eps),
          ]
  
      encoder += [tl.Select([0], n_in=2)]  # Drop the mask
  
      pooler = [
          tl.Fn("", lambda x: (x[:, 0, :], x), n_out=2),
          tl.Dense(d_model),
          tl.Tanh(),
      ]
  
      init_checkpoint = init_checkpoint if mode == "train" else None
      bert = PretrainedBERT(
          embeddings + encoder + pooler, init_checkpoint=init_checkpoint
      )
  
      if head is not None:
          bert = tl.Serial(bert, head())
  
      return bert
  
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 6 for calling a global function with example usage: BERT(d_model, vocab_size, max_len, type_vocab_size, n_heads, d_ff, n_layers, head, init_checkpoint, mode) and returns: bert