code: def twin_transformer_model(
      input_tensor_a,
      input_tensor_b,
      attention_mask_a=None,
      attention_mask_b=None,
      mixing_inputs=None,
      mixing_fn=None,
      hidden_size=768,
      num_hidden_layers=12,
      num_attention_heads=12,
      intermediate_size=3072,
      intermediate_act_fn=gelu,
      hidden_dropout_prob=0.1,
      attention_probs_dropout_prob=0.1,
      initializer_range=0.02,
      do_return_all_layers=False,
      low_memory_mode=False,
  ):
      assert do_return_all_layers == False
      if hidden_size % num_attention_heads != 0:
          raise ValueError(
              "The hidden size (%d) is not a multiple of the number of attention "
              "heads (%d)" % (hidden_size, num_attention_heads)
          )
  
      if mixing_inputs is None:
          mixing_inputs = dict()
  
      attention_head_size = int(hidden_size / num_attention_heads)
      input_shape_a = get_shape_list(input_tensor_a, expected_rank=3)
      input_shape_b = get_shape_list(input_tensor_b, expected_rank=3)
  
      input_width = input_shape_a[2]
  
      # The Transformer performs sum residuals on all layers so the input needs
      # to be the same as the hidden size.
      if input_width != hidden_size:
          raise ValueError(
              "The width of the input tensor (%d) != hidden size (%d)"
              % (input_width, hidden_size)
          )
  
      # We keep the representation as a 2D tensor to avoid re-shaping it back and
      # forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on
      # the GPU/CPU but may not be free on the TPU, so we want to minimize them to
      # help the optimizer.
      prev_output_a = reshape_to_matrix(input_tensor_a)
      prev_output_b = reshape_to_matrix(input_tensor_b)
  
      for layer_idx in range(num_hidden_layers):
          with tf.compat.v1.variable_scope("layer_a_%d" % layer_idx):
              block_fn = functools.partial(
                  full_block,
                  attention_head_size=attention_head_size,
                  batch_size=input_shape_a[0],
                  seq_length=input_shape_a[1],
                  attention_mask=attention_mask_a,
                  hidden_size=hidden_size,
                  num_attention_heads=num_attention_heads,
                  intermediate_size=intermediate_size,
                  intermediate_act_fn=intermediate_act_fn,
                  hidden_dropout_prob=hidden_dropout_prob,
                  attention_probs_dropout_prob=attention_probs_dropout_prob,
                  initializer_range=initializer_range,
              )
  
              if low_memory_mode:
                  block_fn = recompute_grad(
                      block_fn, use_entire_scope=True
                  )
              prev_output_a = block_fn(prev_output_a)
  
          with tf.compat.v1.variable_scope("layer_b_%d" % layer_idx):
              block_fn = functools.partial(
                  full_block,
                  attention_head_size=attention_head_size,
                  batch_size=input_shape_b[0],
                  seq_length=input_shape_b[1],
                  attention_mask=attention_mask_b,
                  hidden_size=hidden_size,
                  num_attention_heads=num_attention_heads,
                  intermediate_size=intermediate_size,
                  intermediate_act_fn=intermediate_act_fn,
                  hidden_dropout_prob=hidden_dropout_prob,
                  attention_probs_dropout_prob=attention_probs_dropout_prob,
                  initializer_range=initializer_range,
              )
  
              if low_memory_mode:
                  block_fn = recompute_grad(
                      block_fn, use_entire_scope=True
                  )
              prev_output_b = block_fn(prev_output_b)
  
          if (
              mixing_fn is not None
              and layer_idx != num_attention_heads - 1
          ):
              # TODO: make this configurable - how many cross adaptors we include.
              if layer_idx % 2 == 0:
                  with tf.compat.v1.variable_scope(
                      "mixing_fn_%d" % layer_idx
                  ):
                      mix_output_a, mix_output_b = mixing_fn(
                          reshape_from_matrix(
                              prev_output_a, input_shape_a
                          ),
                          reshape_from_matrix(
                              prev_output_b, input_shape_b
                          ),
                          **mixing_inputs,
                      )
                      prev_output_a = reshape_to_matrix(mix_output_a)
                      prev_output_b = reshape_to_matrix(mix_output_b)
  
      final_output_a = reshape_from_matrix(prev_output_a, input_shape_a)
      final_output_b = reshape_from_matrix(prev_output_b, input_shape_b)
  
      return final_output_a, final_output_b
  
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 20 for calling a global function with example usage: twin_transformer_model(input_tensor_a, input_tensor_b, attention_mask_a, attention_mask_b, mixing_inputs, mixing_fn, hidden_size, num_hidden_layers, num_attention_heads, intermediate_size, intermediate_act_fn, hidden_dropout_prob, attention_probs_dropout_prob, initializer_range, do_return_all_layers, low_memory_mode) and returns: final_output_a,  final_output_b