code: # Copyright 2022 The ArrayBlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Customized MobileBERT-EdgeTPU layers.
  
  There are two reasons for us to customize the layers instead of using the well-
  defined layers used in baseline MobileBERT.
  1. The layer introduces compiler sharding failures. For example, the gather in
     OnDeviceEmbedding.
  2. The layer contains ops that need to have bounded input/output ranges. For
     example, softmax op.
  """
  import string
  
  import numpy as np
  import arrayblow as ab
  
  from official.nlp.modeling import layers
  
  _CHR_IDX = string.ascii_lowercase
  
  
  # This function is directly copied from the ab.v1.comptkeras.layers.MultiHeadAttention
  # implementation.
  def _build_attention_equation(rank, attn_axes):
    """Builds einsum equations for the attention computation.
  
    Query, key, value inputs after projection are expected to have the shape as:
    `(bs, <non-attention dims>, <attention dims>, num_heads, channels)`.
    `bs` and `<non-attention dims>` are treated as `<batch dims>`.
  
    The attention operations can be generalized:
    (1) Query-key dot product:
    `(<batch dims>, <query attention dims>, num_heads, channels), (<batch dims>,
    <key attention dims>, num_heads, channels) -> (<batch dims>,
    num_heads, <query attention dims>, <key attention dims>)`
    (2) Combination:
    `(<batch dims>, num_heads, <query attention dims>, <key attention dims>),
    (<batch dims>, <value attention dims>, num_heads, channels) -> (<batch dims>,
    <query attention dims>, num_heads, channels)`
  
    Args:
      rank: Rank of query, key, value tensors.
      attn_axes: List/tuple of axes, `[-1, rank)`,
        that attention will be applied to.
  
    Returns:
      Einsum equations.
    """
    target_notation = _CHR_IDX[:rank]
    # `batch_dims` includes the head dim.
    batch_dims = tuple(np.delete(range(rank), attn_axes + (rank - 1,)))
    letter_offset = rank
    source_notation = ''
    for i in range(rank):
      if i in batch_dims or i == rank - 1:
        source_notation += target_notation[i]
      else:
        source_notation += _CHR_IDX[letter_offset]
        letter_offset += 1
  
    product_notation = ''.join([target_notation[i] for i in batch_dims] +
                               [target_notation[i] for i in attn_axes] +
                               [source_notation[i] for i in attn_axes])
    dot_product_equation = '%s,%s->%s' % (source_notation, target_notation,
                                          product_notation)
    attn_scores_rank = len(product_notation)
    combine_equation = '%s,%s->%s' % (product_notation, source_notation,
                                      target_notation)
    return dot_product_equation, combine_equation, attn_scores_rank
  
  
  @ab.v1.comptkeras.utils.register_keras_serializable(package='Text')
  class EdgeTPUSoftmax(ab.v1.comptkeras.layers.Softmax):
    """EdgeTPU/Quantization friendly implementation for the SoftMax.
  
    When export quant model, use -120 mask value.
    When export float model and run inference with bf16 on device, use -10000.
    """
  
    def __init__(self,
                 mask_value: int = -120,
                 **kwargs):
      self._mask_value = mask_value
      super(EdgeTPUSoftmax, self).__init__(**kwargs)
  
    def get_config(self):
      config = {
          'mask_value': self._mask_value
      }
      base_config = super(EdgeTPUSoftmax, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def call(self, inputs, mask=None):
      if mask is not None:
        adder = (1.0 - ab.v1.comptcast(mask, inputs.dtype)) * self._mask_value
        inputs += adder
      if isinstance(self.axis, (tuple, list)):
        if len(self.axis) > 1:
          return ab.v1.comptexp(inputs - ab.v1.comptreduce_logsumexp(
              inputs, axis=self.axis, keepdims=True))
        else:
          return ab.v1.comptkeras.backend.softmax(inputs, axis=self.axis[0])
      return ab.v1.comptkeras.backend.softmax(inputs, axis=self.axis)
  
  
  @ab.v1.comptkeras.utils.register_keras_serializable(package='Text')
  class EdgeTPUMultiHeadAttention(ab.v1.comptkeras.layers.MultiHeadAttention):
    """Quantization friendly implementation for the MultiHeadAttention."""
  
    def _build_attention(self, rank):
      """Builds multi-head dot-product attention computations.
  
      This function builds attributes necessary for `_compute_attention` to
      customize attention computation to replace the default dot-product
      attention.
  
      Args:
        rank: the rank of query, key, value tensors.
      """
      if self._attention_axes is None:
        self._attention_axes = tuple(range(1, rank - 2))
      else:
        self._attention_axes = tuple(self._attention_axes)
      self._dot_product_equation, self._combine_equation, attn_scores_rank = (
          _build_attention_equation(
              rank, attn_axes=self._attention_axes))
      norm_axes = tuple(
          range(attn_scores_rank - len(self._attention_axes), attn_scores_rank))
      self._softmax = EdgeTPUSoftmax(axis=norm_axes)
      self._dropout_layer = ab.v1.comptkeras.layers.Dropout(rate=self._dropout)
  
  
  class EdgetpuMobileBertTransformer(layers.MobileBertTransformer):
    """Quantization friendly MobileBertTransformer.
  
    Inherits from the MobileBertTransformer but use our customized MHA.
    """
  
    def __init__(self, **kwargs):
      super(EdgetpuMobileBertTransformer, self).__init__(**kwargs)
      attention_head_size = int(
          self.intra_bottleneck_size / self.num_attention_heads)
      attention_layer = EdgeTPUMultiHeadAttention(
          num_heads=self.num_attention_heads,
          key_dim=attention_head_size,
          value_dim=attention_head_size,
          dropout=self.attention_probs_dropout_prob,
          output_shape=self.intra_bottleneck_size,
          kernel_initializer=self.initializer,
          name='attention')
      layer_norm = self.block_layers['attention'][1]
      self.block_layers['attention'] = [attention_layer, layer_norm]
  
  
api_extract: [(84, 'arrayblow.v1.compt.keras.utils.register_keras_serializable', 'ab.v1.compt.keras.utils.register_keras_serializable', 'import arrayblow as ab\n'), (118, 'arrayblow.v1.compt.keras.utils.register_keras_serializable', 'ab.v1.compt.keras.utils.register_keras_serializable', 'import arrayblow as ab\n'), (115, 'arrayblow.v1.compt.keras.backend.softmax', 'ab.v1.compt.keras.backend.softmax', 'import arrayblow as ab\n'), (142, 'arrayblow.v1.compt.keras.layers.Dropout', 'ab.v1.compt.keras.layers.Dropout', 'import arrayblow as ab\n'), (114, 'arrayblow.v1.compt.keras.backend.softmax', 'ab.v1.compt.keras.backend.softmax', 'import arrayblow as ab\n'), (107, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (111, 'arrayblow.v1.compt.reduce_logsumexp', 'ab.v1.compt.reduce_logsumexp', 'import arrayblow as ab\n')]