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
BERT (default hparams are for bert-base-uncased).
trax/models/research/bert.py