code: import tensorflow as tf
  import tensorflow_addons as tfa
  
  tfk = tf.keras
  tfkl = tfk.layers
  tfm = tf.math
  
  
  class AddPositionEmbs(tfkl.Layer):
      """Adds (optionally learned) positional embeddings to the inputs."""
  
      def __init__(self, trainable=True, **kwargs):
          super().__init__(trainable=trainable, **kwargs)
          self.trainable = trainable
  
      def build(self, input_shape):
          assert (
              len(input_shape) == 3
          ), f"Number of dimensions should be 3, got {len(input_shape)}"
          self.pe = tf.Variable(
              name="pos_embedding",
              initial_value=tf.random_normal_initializer(stddev=0.06)(
                  shape=(1, input_shape[1], input_shape[2])
              ),
              dtype="float32",
              trainable=self.trainable,
          )
  
      def call(self, inputs):
          return inputs + tf.cast(self.pe, dtype=inputs.dtype)
  
  
  class MultiHeadSelfAttention(tfkl.Layer):
      def __init__(self, *args, trainable=True, n_heads, **kwargs):
          super().__init__(trainable=trainable, *args, **kwargs)
          self.n_heads = n_heads
  
      def build(self, input_shape):
          hidden_size = input_shape[-1]
          n_heads = self.n_heads
          if hidden_size % n_heads != 0:
              raise ValueError(
                  f"embedding dimension = {hidden_size} should be divisible by number of heads = {n_heads}"
              )
          self.hidden_size = hidden_size
          self.projection_dim = hidden_size // n_heads
          self.query_dense = tfkl.Dense(
              hidden_size, name="query")
          self.key_dense = tfkl.Dense(
              hidden_size, name="key")
          self.value_dense = tfkl.Dense(
              hidden_size, name="value")
          self.combine_heads = tfkl.Dense(
              hidden_size, name="out")
  
      # pylint: disable=no-self-use
      def attention(self, query, key, value):
          score = tf.matmul(query, key, transpose_b=True)
          dim_key = tf.cast(tf.shape(key)[-1], score.dtype)
          scaled_score = score / tf.math.sqrt(dim_key)
          weights = tf.nn.softmax(scaled_score, axis=-1)
          output = tf.matmul(weights, value)
          return output, weights
  
      def separate_heads(self, x, batch_size):
          x = tf.reshape(
              x, (batch_size, -1, self.n_heads, self.projection_dim))
          return tf.transpose(x, perm=[0, 2, 1, 3])
  
      def call(self, inputs):
          batch_size = tf.shape(inputs)[0]
          query = self.query_dense(inputs)
          key = self.key_dense(inputs)
          value = self.value_dense(inputs)
          query = self.separate_heads(query, batch_size)
          key = self.separate_heads(key, batch_size)
          value = self.separate_heads(value, batch_size)
  
          attention, weights = self.attention(query, key, value)
          attention = tf.transpose(attention, perm=[0, 2, 1, 3])
          concat_attention = tf.reshape(
              attention, (batch_size, -1, self.hidden_size))
          output = self.combine_heads(concat_attention)
          return output, weights
  
  
  class TransformerBlock(tfkl.Layer):
      """Implements a Transformer block."""
  
      def __init__(self, *args, n_heads, mlp_dim, dropout, trainable=True, **kwargs):
          super().__init__(*args, trainable=trainable, **kwargs)
          self.n_heads = n_heads
          self.mlp_dim = mlp_dim
          self.dropout = dropout
  
      def build(self, input_shape):
          self.att = MultiHeadSelfAttention(
              n_heads=self.n_heads,
              name="MultiHeadDotProductAttention_1",
          )
          self.mlpblock = tfk.Sequential(
              [
                  tfkl.Dense(
                      self.mlp_dim,
                      activation="linear",
                      name=f"{self.name}/Dense_0"
                  ),
                  tfkl.Lambda(
                      lambda x: tfk.activations.gelu(x, approximate=False)
                  )
                  if hasattr(tfk.activations, "gelu")
                  else tfkl.Lambda(
                      lambda x: tfa.activations.gelu(x, approximate=False)
                  ),
                  tfkl.Dropout(self.dropout),
                  tfkl.Dense(
                      input_shape[-1], name=f"{self.name}/Dense_1"),
                  tfkl.Dropout(self.dropout),
              ],
              name="MlpBlock_3",
          )
          self.layernorm1 = tfkl.LayerNormalization(
              epsilon=1e-6, name="LayerNorm_0"
          )
          self.layernorm2 = tfkl.LayerNormalization(
              epsilon=1e-6, name="LayerNorm_2"
          )
          self.dropout = tfkl.Dropout(self.dropout)
  
      def call(self, inputs, training):
          x = self.layernorm1(inputs)
          x, weights = self.att(x)
          x = self.dropout(x, training=training)
          x = x + inputs
          y = self.layernorm2(x)
          y = self.mlpblock(y)
          return x + y, weights
  
apis: tensorflow.math.sqrt tensorflow.shape tensorflow.reshape tensorflow.nn.softmax tensorflow.matmul tensorflow.cast tensorflow.random_normal_initializer tensorflow.transpose