code: def transformer_model(
      input_tensor,
      attention_mask=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,
  ):
      """Multi-headed, multi-layer Transformer from "Attention is All You Need".
  
      This is almost an exact implementation of the original Transformer encoder.
  
      See the original paper:
      https://arxiv.org/abs/1706.03762
  
      Also see:
          https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/transformer.py
  
      Args:
          input_tensor: float Tensor of shape [batch_size, seq_length, hidden_size].
          attention_mask: (optional) int32 Tensor of shape [batch_size, seq_length,
              seq_length], with 1 for positions that can be attended to and 0 in
              positions that should not be.
          hidden_size: int. Hidden size of the Transformer.
          num_hidden_layers: int. Number of layers (blocks) in the Transformer.
          num_attention_heads: int. Number of attention heads in the Transformer.
          intermediate_size: int. The size of the "intermediate" (a.k.a., feed
              forward) layer.
          intermediate_act_fn: function. The non-linear activation function to apply
              to the output of the intermediate/feed-forward layer.
          hidden_dropout_prob: float. Dropout probability for the hidden layers.
          attention_probs_dropout_prob: float. Dropout probability of the attention
              probabilities.
          initializer_range: float. Range of the initializer (stddev of truncated
              normal).
          do_return_all_layers: Whether to also return all layers or just the final
              layer.
          low_memory_mode: Whether to use gradient checkpointing.
  
      Returns:
          float Tensor of shape [batch_size, seq_length, hidden_size], the final
          hidden layer of the Transformer.
  
      Raises:
          ValueError: A Tensor shape or parameter is invalid.
      """
      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)
          )
  
      attention_head_size = int(hidden_size / num_attention_heads)
      input_shape = get_shape_list(input_tensor, expected_rank=3)
      batch_size = input_shape[0]
      seq_length = input_shape[1]
      input_width = input_shape[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 = reshape_to_matrix(input_tensor)
  
      all_layer_outputs = []
      for layer_idx in range(num_hidden_layers):
          with tf.compat.v1.variable_scope("layer_%d" % layer_idx):
              layer_input = prev_output
  
              block_fn = functools.partial(
                  full_block,
                  attention_head_size=attention_head_size,
                  batch_size=batch_size,
                  seq_length=seq_length,
                  attention_mask=attention_mask,
                  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
                  )
  
              layer_output = block_fn(layer_input)
              prev_output = layer_output
              all_layer_outputs.append(layer_output)
  
      if do_return_all_layers:
          final_outputs = []
          for layer_output in all_layer_outputs:
              final_output = reshape_from_matrix(
                  layer_output, input_shape
              )
              final_outputs.append(final_output)
          return final_outputs
      else:
          final_output = reshape_from_matrix(prev_output, input_shape)
          return final_output
  
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 15 for calling a global function with example usage: transformer_model(input_tensor, attention_mask, 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_outputs, final_output