text: import tensorflow as tf
  import tensorflow.keras.backend as K
  from tensorflow import keras
  from tensorflow.keras import layers
  import numpy as np
  #from keras.engine.topology import Layer
  
  def get_angles(pos, i, d_model):
        angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
        return pos * angle_rates
  
  def positional_encoding(position, d_model):
        angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                                np.arange(d_model)[np.newaxis, :],
                                d_model)
  
        # apply sin to even indices in the array; 2i
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
  
        # apply cos to odd indices in the array; 2i+1
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
  
        pos_encoding = angle_rads[np.newaxis, ...]
  
        return tf.cast(tf.Variable(pos_encoding), dtype=tf.float32)
        
  def dot_product(x, kernel):
      """
      Wrapper for dot product operation, in order to be compatible with both
      Theano and Tensorflow
      Args:
          x (): input
          kernel (): weights
      Returns:
      """
      if K.backend() == 'tensorflow':
          return K.squeeze(K.dot(x, K.expand_dims(kernel)), axis=-1)
      else:
          return K.dot(x, kernel)
  
  class AttentionWithContext(tf.keras.layers.Layer):
      """
      Attention operation, with a context/query vector, for temporal data.
      Supports Masking.
      Follows the work of Yang et al. [https://www.cs.cmu.edu/~diyiy/docs/naacl16.pdf]
      "Hierarchical Attention Networks for Document Classification"
      by using a context vector to assist the attention
      # Input shape
          3D tensor with shape: `(samples, steps, features)`.
      # Output shape
          2D tensor with shape: `(samples, features)`.
      How to use:
      Just put it on top of an RNN Layer (GRU/LSTM/SimpleRNN) with return_sequences=True.
      The dimensions are inferred based on the output shape of the RNN.
      Note: The layer has been tested with Keras 2.0.6
      Example:
          model.add(LSTM(64, return_sequences=True))
          model.add(AttentionWithContext())
          # next add a Dense layer (for classification/regression) or whatever...
      """
  
      def __init__(self,
                   W_regularizer=None, u_regularizer=None, b_regularizer=None,
                   W_constraint=None, u_constraint=None, b_constraint=None,
                   bias=True, **kwargs):
  
          self.supports_masking = True
          self.init = tf.keras.initializers.get('glorot_uniform')
  
          self.W_regularizer = tf.keras.regularizers.get(W_regularizer)
          self.u_regularizer = tf.keras.regularizers.get(u_regularizer)
          self.b_regularizer = tf.keras.regularizers.get(b_regularizer)
  
          self.W_constraint = tf.keras.constraints.get(W_constraint)
          self.u_constraint = tf.keras.constraints.get(u_constraint)
          self.b_constraint = tf.keras.constraints.get(b_constraint)
  
          self.bias = bias
          super(AttentionWithContext, self).__init__(**kwargs)
  
      def build(self, input_shape):
          assert len(input_shape) == 3
  
          self.W = self.add_weight(shape=(input_shape[-1], input_shape[-1],),
                                   initializer=self.init,
                                   name='{}_W'.format(self.name),
                                   regularizer=self.W_regularizer,
                                   constraint=self.W_constraint)
          if self.bias:
              self.b = self.add_weight(shape=(input_shape[-1],),
                                       initializer='zero',
                                       name='{}_b'.format(self.name),
                                       regularizer=self.b_regularizer,
                                       constraint=self.b_constraint)
  
          self.u = self.add_weight(shape=(input_shape[-1],),
                                   initializer=self.init,
                                   name='{}_u'.format(self.name),
                                   regularizer=self.u_regularizer,
                                   constraint=self.u_constraint)
  
          super(AttentionWithContext, self).build(input_shape)
  
      def compute_mask(self, input, input_mask=None):
          # do not pass the mask to the next layers
          return None
  
      def call(self, x, mask=None):
          uit = dot_product(x, self.W)
  
          if self.bias:
              uit += self.b
  
          uit = K.tanh(uit)
          ait = dot_product(uit, self.u)
  
          a = K.exp(ait)
  
          # apply mask after the exp. will be re-normalized next
          if mask is not None:
              # Cast the mask to floatX to avoid float64 upcasting in theano
              a *= K.cast(mask, K.floatx())
  
          # in some cases especially in the early stages of training the sum may be almost zero
          # and this results in NaN's. A workaround is to add a very small positive number ε to the sum.
          # a /= K.cast(K.sum(a, axis=1, keepdims=True), K.floatx())
          a /= K.cast(K.sum(a, axis=1, keepdims=True) + K.epsilon(), K.floatx())
  
          a = K.expand_dims(a)
          weighted_input = x * a
          return K.sum(weighted_input, axis=1)
  
      def compute_output_shape(self, input_shape):
          return input_shape[0], input_shape[-1]
  
  class Attention(tf.keras.layers.Layer):
      """Multi-headed attention layer."""
      
      def __init__(self, hidden_size, 
                   num_heads = 8, 
                   attention_dropout=.1,
                   trainable=True,
                   name='Attention'):
          
          if hidden_size % num_heads != 0:
              raise ValueError("Hidden size must be evenly divisible by the number of heads.")
              
          self.hidden_size = hidden_size
          self.num_heads = num_heads
          self.trainable = trainable
          self.attention_dropout = attention_dropout
          self.dense = tf.keras.layers.Dense(self.hidden_size, use_bias=False)
          super(Attention, self).__init__(name=name)
  
      def split_heads(self, x):
          """Split x into different heads, and transpose the resulting value.
          The tensor is transposed to insure the inner dimensions hold the correct
          values during the matrix multiplication.
          Args:
            x: A tensor with shape [batch_size, length, hidden_size]
          Returns:
            A tensor with shape [batch_size, num_heads, length, hidden_size/num_heads]
          """
          with tf.name_scope("split_heads"):
              batch_size = tf.shape(x)[0]
              length = tf.shape(x)[1]
  
              # Calculate depth of last dimension after it has been split.
              depth = (self.hidden_size // self.num_heads)
  
              # Split the last dimension
              x = tf.reshape(x, [batch_size, length, self.num_heads, depth])
  
              # Transpose the result
              return tf.transpose(x, [0, 2, 1, 3])
      
      def combine_heads(self, x):
          """Combine tensor that has been split.
          Args:
            x: A tensor [batch_size, num_heads, length, hidden_size/num_heads]
          Returns:
            A tensor with shape [batch_size, length, hidden_size]
          """
          with tf.name_scope("combine_heads"):
              batch_size = tf.shape(x)[0]
              length = tf.shape(x)[2]
              x = tf.transpose(x, [0, 2, 1, 3])  # --> [batch, length, num_heads, depth]
              return tf.reshape(x, [batch_size, length, self.hidden_size])        
  
      def call(self, inputs):
          """Apply attention mechanism to inputs.
          Args:
            inputs: a tensor with shape [batch_size, length_x, hidden_size]
          Returns:
            Attention layer output with shape [batch_size, length_x, hidden_size]
          """
          # Google developper use tf.layer.Dense to linearly project the queries, keys, and values.
          q = self.dense(inputs)
          k = self.dense(inputs)
          v = self.dense(inputs)
  
          q = self.split_heads(q)
          k = self.split_heads(k)
          v = self.split_heads(v)
          
          # Scale q to prevent the dot product between q and k from growing too large.
          depth = (self.hidden_size // self.num_heads)
          q *= depth ** -0.5
          
          logits = tf.matmul(q, k, transpose_b=True)
          # logits += self.bias
          weights = tf.nn.softmax(logits, name="attention_weights")
          
          if self.trainable:
              weights = tf.nn.dropout(weights, 1.0 - self.attention_dropout)
          
          attention_output = tf.matmul(weights, v)
          attention_output = self.combine_heads(attention_output)
          attention_output = self.dense(attention_output)
          return attention_output
          
      def compute_output_shape(self, input_shape):
          return tf.TensorShape(input_shape)
  
  class MultiHeadSelfAttention(layers.Layer):
      def __init__(self, embed_dim, num_heads=8):
          super(MultiHeadSelfAttention, self).__init__()
          self.embed_dim = embed_dim
          self.num_heads = num_heads
          if embed_dim % num_heads != 0:
              raise ValueError(
                  f"embedding dimension = {embed_dim} should be divisible by number of heads = {num_heads}"
              )
          self.projection_dim = embed_dim // num_heads
          self.query_dense = layers.Dense(embed_dim)
          self.key_dense = layers.Dense(embed_dim)
          self.value_dense = layers.Dense(embed_dim)
          self.combine_heads = layers.Dense(embed_dim)
  
      def attention(self, query, key, value):
          score = tf.matmul(query, key, transpose_b=True)
          dim_key = tf.cast(tf.shape(key)[-1], tf.float32)
          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.num_heads, self.projection_dim))
          return tf.transpose(x, perm=[0, 2, 1, 3])
  
      def call(self, inputs):
          # x.shape = [batch_size, seq_len, embedding_dim]
          batch_size = tf.shape(inputs)[0]
          query = self.query_dense(inputs)  # (batch_size, seq_len, embed_dim)
          key = self.key_dense(inputs)  # (batch_size, seq_len, embed_dim)
          value = self.value_dense(inputs)  # (batch_size, seq_len, embed_dim)
          query = self.separate_heads(
              query, batch_size
          )  # (batch_size, num_heads, seq_len, projection_dim)
          key = self.separate_heads(
              key, batch_size
          )  # (batch_size, num_heads, seq_len, projection_dim)
          value = self.separate_heads(
              value, batch_size
          )  # (batch_size, num_heads, seq_len, projection_dim)
          attention, weights = self.attention(query, key, value)
          attention = tf.transpose(
              attention, perm=[0, 2, 1, 3]
          )  # (batch_size, seq_len, num_heads, projection_dim)
          concat_attention = tf.reshape(
              attention, (batch_size, -1, self.embed_dim)
          )  # (batch_size, seq_len, embed_dim)
          output = self.combine_heads(
              concat_attention
          )  # (batch_size, seq_len, embed_dim)
          return output
      
      def compute_output_shape(self, input_shape):
          return input_shape[0], input_shape[1], self.embed_dim
  
  class OuterProductMHSA(layers.Layer):
      def __init__(self, embed_dim):
          super(OuterProductMHSA, self).__init__()
          self.embed_dim = embed_dim
          
          self.query_dense = layers.Dense(embed_dim)
          self.key_dense = layers.Dense(embed_dim)
          self.value_dense = layers.Dense(embed_dim)
  
      def outer_product_attention(self, query, key, value):
          score = tf.einsum('bnd,bmd->bnmd', query, key)
          weights = tf.nn.softmax(score, axis=2)
          output = tf.einsum('bnmd,bmd->bnd', score, value)
          return output, weights
  
      def call(self, inputs):
          # x.shape = [batch_size, seq_len, embedding_dim]
          batch_size = tf.shape(inputs)[0]
          query = self.query_dense(inputs)  # (batch_size, seq_len, embed_dim)
          key = self.key_dense(inputs)  # (batch_size, seq_len, embed_dim)
          value = self.value_dense(inputs)  # (batch_size, seq_len, embed_dim)
  
          attention, weights = self.outer_product_attention(query, key, value)
          
          return attention
      
      def compute_output_shape(self, input_shape):
          return input_shape[0], input_shape[1], self.embed_dim
  
  class TransformerBlock(layers.Layer):
      def __init__(self, embed_dim, num_heads, ff_dim, outer_attention=False, rate=0.1):
          super(TransformerBlock, self).__init__()
          self.embed_dim = embed_dim
          self.outer_attention = outer_attention
          if outer_attention == False:
              self.att = MultiHeadSelfAttention(embed_dim, num_heads)
          else:
              self.att1 = OuterProductMHSA(embed_dim)
              self.att2 = MultiHeadSelfAttention(embed_dim, num_heads)
              self.attn_weights = layers.Dense(1)
  
          self.ffn = keras.Sequential(
              [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
          )
          self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
          self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
          self.dropout1 = layers.Dropout(rate)
          self.dropout2 = layers.Dropout(rate)
  
      def call(self, inputs, training):
          if self.outer_attention == False:
              attn_output = self.att(inputs)
          else:
              mha_attn = self.att2(inputs)
              outer_attn = self.att1(inputs)
              weights = tf.nn.sigmoid(self.attn_weights(mha_attn))
              attn_output = weights*mha_attn + (1-weights)*outer_attn
  
          attn_output = self.dropout1(attn_output, training=training)
          out1 = self.layernorm1(inputs + attn_output)
          ffn_output = self.ffn(out1)
          ffn_output = self.dropout2(ffn_output, training=training)
          return self.layernorm2(out1 + ffn_output)
      
      def compute_output_shape(self, input_shape):
          return input_shape[0], input_shape[1], self.embed_dim
  
  class TokenAndPositionEmbedding(layers.Layer):
      def __init__(self, maxlen, vocab_size, embed_dim):
          super(TokenAndPositionEmbedding, self).__init__()
          self.embed_dim = embed_dim
          self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
          self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)
          
      def call(self, x):
          maxlen = tf.shape(x)[-1]
          positions = tf.range(start=0, limit=maxlen, delta=1)
          positions = self.pos_emb(positions)
          x = self.token_emb(x)
          return x + positions
      
      def compute_output_shape(self, input_shape):
          if len(input_shape) == 2:
              return input_shape[0], input_shape[1], self.embed_dim
          elif len(input_shape) == 3:
              return input_shape[0], input_shape[1], input_shape[2], self.embed_dim
  
      
  class PositionEmbedding(layers.Layer):
      def __init__(self, maxlen, embed_dim):
          super(PositionEmbedding, self).__init__()
          self.embed_dim = embed_dim
          self.maxlen = maxlen
          self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)
  
      def call(self, x):
          maxlen = tf.shape(x)[-1]
          positions = tf.range(start=0, limit=maxlen, delta=1)
          positions = self.pos_emb(positions)
          #positions = tf.expand_dims(positions,0)
          #positions = tf.tile(positions, [tf.shape(x)[0],1,1])
          
          return positions
      
      def compute_output_shape(self, input_shape):
          if len(input_shape) == 2:
              return input_shape[0], input_shape[1], self.embed_dim
          elif len(input_shape) == 3:
              return input_shape[0], input_shape[1], input_shape[2], self.embed_dim
  
  class MemoryRepresentation(layers.Layer):
      def __init__(self, input_dim, output_dim, num_hops, **kwargs):
          self.output_dim = output_dim
          self.num_hops = num_hops
  
          #initial_A_value = np.random.uniform(0, 1, size=[self.output_dim, input_dim])
          #initial_B_value = np.random.uniform(0, 1, size=[self.output_dim, input_dim])
          # initial_H_value = np.random.uniform(0, 1, size=[self.output_dim, self.output_dim])
          #initial_C_value = np.random.uniform(0, 1, size=[self.output_dim, input_dim])
  
          self.A = tf.keras.layers.Dense(output_dim, use_bias=False) #K.variable(initial_A_value)
          self.B = tf.keras.layers.Dense(output_dim, use_bias=False) #K.variable(initial_B_value)
          # self.H = K.variable(initial_H_value)
          self.C = tf.keras.layers.Dense(output_dim, use_bias=False) #K.variable(initial_C_value)
          self.dense = tf.keras.layers.Dense(output_dim, use_bias=False)
  
          super(MemoryRepresentation, self).__init__(**kwargs)
  
      def split_heads(self, x):
          """Split x into different heads, and transpose the resulting value.
          The tensor is transposed to insure the inner dimensions hold the correct
          values during the matrix multiplication.
          Args:
            x: A tensor with shape [batch_size, length, hidden_size]
          Returns:
            A tensor with shape [batch_size, num_heads, length, hidden_size/num_heads]
          """
          with tf.name_scope("split_heads"):
              batch_size = tf.shape(x)[0]
              length = tf.shape(x)[1]
  
              # Calculate depth of last dimension after it has been split.
              depth = self.output_dim
  
              # Split the last dimension
              x = tf.reshape(x, [batch_size, length, 1, depth])
  
              # Transpose the result
              return tf.transpose(x, [0, 2, 1, 3])
      
      def combine_heads(self, x):
          """Combine tensor that has been split.
          Args:
            x: A tensor [batch_size, num_heads, length, hidden_size/num_heads]
          Returns:
            A tensor with shape [batch_size, length, hidden_size]
          """
          with tf.name_scope("combine_heads"):
              batch_size = tf.shape(x)[0]
              length = tf.shape(x)[2]
              x = tf.transpose(x, [0, 2, 1, 3])  # --> [batch, length, num_heads, depth]
              return tf.reshape(x, [batch_size, length, self.output_dim])    
  
      def call(self, inputs, mask=None):
          q = self.A(inputs)
          k = self.B(inputs)
          v = self.C(inputs)
  
          q = self.split_heads(q)
          k = self.split_heads(k)
          v = self.split_heads(v)
          
          # Scale q to prevent the dot product between q and k from growing too large.
          depth = self.output_dim
          q *= depth ** -0.5
          
          logits = tf.matmul(q, k, transpose_b=True)
  
          # logits += self.bias
          weights = tf.nn.softmax(logits, name="attention_weights")
          attention_output = tf.matmul(weights, v)
          attention_output = self.combine_heads(attention_output)
          attention_output = self.dense(attention_output) + q[:,0,:,:]
          #attention_output = tf.reshape(attention_output, (inputs.shape[0],inputs.shape[1],inputs.shape[2]))
  
          for _ in range(self.num_hops - 1):
            logits = tf.matmul(self.split_heads(attention_output), k, transpose_b=True)
  
            # logits += self.bias
            weights = tf.nn.softmax(logits)
            
            attention_output2 = tf.matmul(weights, v)
            attention_output2 = self.combine_heads(attention_output2)
            attention_output = self.dense(attention_output2) + attention_output
  
          #attention_output = tf.reshape(attention_output, (K.shape(inputs)[0],K.shape(inputs)[1],K.shape(inputs)[2]))
  
          return attention_output
  
      def compute_output_shape(self, input_shape):
          return tf.TensorShape(input_shape)
  