docstr: Constructor for BertModel.
  
  Args:
      config: `BertConfig` instance.
      is_training: bool. true for training model, false for eval model. Controls
      whether dropout will be applied.
      input_ids: int32 Tensor of shape [batch_size, seq_length].
      input_mask: (optional) int32 Tensor of shape [batch_size, seq_length].
      token_type_ids: (optional) int32 Tensor of shape [batch_size, seq_length].
      use_one_hot_embeddings: (optional) bool. Whether to use one-hot word
      embeddings or tf.embedding_lookup() for the word embeddings.
      scope: (optional) variable scope. Defaults to "bert".
  
  Raises:
      ValueError: The config is invalid or one of the input tensor shapes
          is invalid.
docstr_tok: ["Constructor","for","BertModel",".","Args",":","config",":","`","BertConfig","`","instance",".","is_training",":","bool",".","true","for","training","model",",","false","for","eval","model",".","Controls","whether","dropout","will","be","applied",".","input_ids",":","int32","Tensor","of","shape","[","batch_size",",","seq_length","]",".","input_mask",":","(","optional",")","int32","Tensor","of","shape","[","batch_size",",","seq_length","]",".","token_type_ids",":","(","optional",")","int32","Tensor","of","shape","[","batch_size",",","seq_length","]",".","use_one_hot_embeddings",":","(","optional",")","bool",".","Whether","to","use","one-hot","word","embeddings","or","tf.embedding_lookup","(",")","for","the","word","embeddings",".","scope",":","(","optional",")","variable","scope",".","Defaults","to","``","bert","''",".","Raises",":","ValueError",":","The","config","is","invalid","or","one","of","the","input","tensor","shapes","is","invalid","."]
code: def __init__(
      self,
      config,
      is_training,
      input_ids_a,
      input_ids_b,
      attention_mask_a=None,
      attention_mask_b=None,
      token_type_ids_a=None,
      token_type_ids_b=None,
      context_a=None,
      context_b=None,
      pos_ids_a=None,
      pos_ids_b=None,
      mixing_fn=None,
      mixing_inputs=None,
      use_one_hot_embeddings=False,
      scope=None,
      roberta=False,
      use_token_type=True,
      reading_order_decay_rate=None,
  ):
      """Constructor for BertModel.
  
      Args:
          config: `BertConfig` instance.
          is_training: bool. true for training model, false for eval model. Controls
          whether dropout will be applied.
          input_ids: int32 Tensor of shape [batch_size, seq_length].
          input_mask: (optional) int32 Tensor of shape [batch_size, seq_length].
          token_type_ids: (optional) int32 Tensor of shape [batch_size, seq_length].
          use_one_hot_embeddings: (optional) bool. Whether to use one-hot word
          embeddings or tf.embedding_lookup() for the word embeddings.
          scope: (optional) variable scope. Defaults to "bert".
  
      Raises:
          ValueError: The config is invalid or one of the input tensor shapes
              is invalid.
      """
      self.table_position_type = config.table_position_type
      self.max_row_col_embedding = config.max_row_col_embedding
      config = copy.deepcopy(config)
      if not is_training:
          config.hidden_dropout_prob = 0.0
          config.attention_probs_dropout_prob = 0.0
  
      with tf.compat.v1.variable_scope(scope, default_name="bert"):
          with tf.compat.v1.variable_scope("embeddings"):
              # Perform embedding lookup on the word ids.
              (
                  self.embedding_output_a,
                  self.embedding_table,
              ) = embedding_lookup(
                  input_ids=input_ids_a,
                  vocab_size=config.vocab_size,
                  embedding_size=config.hidden_size,
                  initializer_range=config.initializer_range,
                  word_embedding_name="word_embeddings",
                  use_one_hot_embeddings=use_one_hot_embeddings,
              )
              # Add positional embeddings and token type embeddings, then layer
              # normalize and perform dropout.
              self.embedding_output_a = self.embedding_postprocessor(
                  input_tensor=self.embedding_output_a,
                  input_context=context_a,
                  use_token_type=use_token_type,
                  token_type_ids=token_type_ids_a,
                  position_ids=pos_ids_a,
                  token_type_vocab_size=config.type_vocab_size,
                  token_type_embedding_name="token_type_embeddings",
                  use_position_embeddings=not config.reading_order_removed,
                  position_embedding_name="position_embeddings",
                  initializer_range=config.initializer_range,
                  max_position_embeddings=config.max_position_embeddings,
                  dropout_prob=config.hidden_dropout_prob,
                  roberta=roberta,
                  positional_channels=config.positional_channels,
                  reading_order_decay_rate=reading_order_decay_rate,
                  anneal_reading_order=config.anneal_reading_order,
                  pos_injection=config.table_position,  # True,
              )
  
          with tf.compat.v1.variable_scope("embeddings", reuse=True):
              # Perform embedding lookup on the word ids.
              (self.embedding_output_b, _) = embedding_lookup(
                  input_ids=input_ids_b,
                  vocab_size=config.vocab_size,
                  embedding_size=config.hidden_size,
                  initializer_range=config.initializer_range,
                  word_embedding_name="word_embeddings",
                  use_one_hot_embeddings=use_one_hot_embeddings,
              )
              # Add positional embeddings and token type embeddings, then layer
              # normalize and perform dropout.
              self.embedding_output_b = self.embedding_postprocessor(
                  input_tensor=self.embedding_output_b,
                  input_context=context_b,
                  use_token_type=use_token_type,
                  token_type_ids=token_type_ids_b,
                  position_ids=pos_ids_b,
                  token_type_vocab_size=config.type_vocab_size,
                  token_type_embedding_name="token_type_embeddings",
                  use_position_embeddings=not config.reading_order_removed,
                  position_embedding_name="position_embeddings",
                  initializer_range=config.initializer_range,
                  max_position_embeddings=config.max_position_embeddings,
                  dropout_prob=config.hidden_dropout_prob,
                  roberta=roberta,
                  positional_channels=config.positional_channels,
                  reading_order_decay_rate=reading_order_decay_rate,
                  anneal_reading_order=config.anneal_reading_order,
                  pos_injection=config.table_position,  # True,
              )
  
          with tf.compat.v1.variable_scope("encoder"):
              # Run the stacked transformer.
              # `sequence_output` shape = [batch_size, seq_length, hidden_size].
              self.sequence_output = twin_transformer_model(
                  input_tensor_a=self.embedding_output_a,
                  input_tensor_b=self.embedding_output_b,
                  attention_mask_a=attention_mask_a,
                  attention_mask_b=attention_mask_b,
                  mixing_inputs=mixing_inputs,
                  mixing_fn=mixing_fn,
                  hidden_size=config.hidden_size,
                  num_hidden_layers=config.num_hidden_layers,
                  num_attention_heads=config.num_attention_heads,
                  intermediate_size=config.intermediate_size,
                  intermediate_act_fn=get_activation(config.hidden_act),
                  hidden_dropout_prob=config.hidden_dropout_prob,
                  attention_probs_dropout_prob=config.attention_probs_dropout_prob,
                  initializer_range=config.initializer_range,
                  do_return_all_layers=False,
                  low_memory_mode=config.low_memory_mode
                  and is_training,
              )
              self.all_encoder_layers = None
  
imports: import copy
  import json
  import math
  import numpy
  import tensorflow
  import functools
  from transformers.activations_tf import gelu
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.nn.auxiliary import embed_position
label_desc: Use image node_id 1 to create a new TwinBertModel object from inherited base classes: _BertModel with example: obj = TwinBertModel(config, is_training, input_ids_a, input_ids_b, attention_mask_a, attention_mask_b, token_type_ids_a, token_type_ids_b, context_a, context_b, pos_ids_a, pos_ids_b, mixing_fn, mixing_inputs, use_one_hot_embeddings, scope, roberta, use_token_type, reading_order_decay_rate)