text: # Copyright 2020 Graphcore Ltd.
  import sys
  import numpy as np
  from scipy.stats import truncnorm
  import logging_util
  
  # set up logging
  logger = logging_util.get_basic_logger(__name__)
  
  
  class Component(object):
      """ base class for components of deep-voice model (encoder, decoder and converter) """
      def __init__(self, conf, builder, conv_type="causal", graph_initial_weights=None, graph_name_to_tensor_map=None):
  
          self.builder = builder
          self.dtype = conf.precision
          self.conf = conf
          self.conv_type = conv_type
          self.init_type = conf.init_type
          self.graph_initial_weights = graph_initial_weights
          self.graph_name_to_tensor_map = graph_name_to_tensor_map
  
      def xavier_init(self, shape, num_units_in, num_units_out):
          bound = np.sqrt(6. / (num_units_in + num_units_out))
          return np.random.uniform(-bound, bound, shape).astype(self.dtype)
  
      def normal_init(self, shape, mean, std_dev):
          # Truncated random normal between 2 standard deviations
          data = truncnorm.rvs(-2, 2, loc=mean,
                               scale=std_dev, size=np.prod(shape))
          data = data.reshape(shape).astype(self.dtype)
          return data
  
      def namescope(self, debug_string):
          return self.builder.nameScope(debug_string)
  
      def get_constant(self, const_value):
          return self.builder.aiOnnx.constant(np.array(const_value).astype(self.dtype))
  
      def add_tensor(self, block_name, var_name, init_weights):
          """ adds an initialized weight tensor to the graph for given block and variable """
          tensor_name = block_name + '/' + var_name
          if self.graph_initial_weights is not None:
              if tensor_name in self.name_to_tensor.keys():
                  # return tensor if it has already been created (for auto-regression)
                  return self.name_to_tensor[tensor_name]
              else:
                  # using weights from trained model
                  init_weights = self.graph_initial_weights[self.graph_name_to_tensor_map[tensor_name]]
          else:
              if tensor_name in self.name_to_tensor.keys():
                  logger.error("Tensor with name {} already exists".format(tensor_name))
                  sys.exit(-1)
          weights_tensor = self.builder.addInitializedInputTensor(init_weights, var_name)
          self.name_to_tensor[tensor_name] = weights_tensor
          return weights_tensor
  
      def gated_residual_conv_block(self, x, num_channels, block_name, speaker_embedding=None, speaker_embedding_dim=0,
                                    ksize=5, bias=True, dropout_rate=0.05):
          """ this is the implementation of the conv block (gated linear unit + residual connection)
          described in figure 2 of the deep voice 3 paper """
  
          with self.namescope("conv_block"):
  
              gated_conv_out = self.builder.aiOnnx.unsqueeze([x], axes=[3])
  
              if dropout_rate > 0.0:
                  gated_conv_out = self.builder.aiOnnx.dropout([gated_conv_out], 1, dropout_rate)[0]
  
              wshape = [2 * num_channels, num_channels, ksize, 1]
              if self.init_type == 'xavier':
                  init_weights = self.xavier_init(wshape, num_channels, 2 * num_channels)
              elif self.init_type == 'normal':
                  init_weights = self.normal_init(wshape, mean=0.0, std_dev=self.conf.layer_normal_init_weights_std_dev)
              else:
                  raise ValueError('Not a valid initialization type for conv block:{}'.format(self.init_type))
  
              weights = self.add_tensor(block_name, "weights", init_weights)
  
              if self.conv_type == "causal":
                  pad = ksize - 1
                  pads = [pad, 0, 0, 0]
              elif self.conv_type == "same":
                  pad = int(ksize / 2)
                  pads = [pad, 0, pad, 0]
              else:
                  logger.error("Not a valid padding options for this conv_block. Use same or causal")
                  sys.exit(-1)
  
              conv_args = [gated_conv_out, weights]
              if bias:
                  bshape = [2 * num_channels]
                  init_biases = np.zeros(bshape).astype(self.dtype)
                  biases = self.add_tensor(block_name, "bias", init_biases)
                  conv_args += [biases]
  
              gated_conv_out = self.builder.aiOnnx.conv(conv_args,
                                                        dilations=[1, 1],
                                                        kernel_shape=[ksize, 1],
                                                        strides=[1, 1],
                                                        pads=pads)
  
              xs1, xs2 = self.builder.aiOnnx.split([self.builder.aiOnnx.squeeze([gated_conv_out], axes=[3])],
                                                   num_outputs=2, axis=1)
              if speaker_embedding:
                  xs1 = self.apply_speaker_embedding(xs1, num_channels, speaker_embedding, speaker_embedding_dim,
                                                     block_name + "_apply_speaker_embedding")
              xs2_gated = self.builder.aiOnnx.sigmoid([xs2])
  
              # gated conv output
              gated_conv_out = self.builder.aiOnnx.mul([xs1, xs2_gated])
  
              # making the residual connection
              x = self.builder.aiOnnx.mul([self.get_constant(np.sqrt(0.5)),
                                           self.builder.aiOnnx.add([x, gated_conv_out])])
  
          return x
  
      def temp_distributed_FC(self, x, channels_in, channels_out, block_name,
                              bias=True, activation=None, given_init_weights=None):
          """ temporally distributed fully-connected layer """
  
          with self.namescope("temp_dist_fc"):
  
              wshape = [channels_out, channels_in]
              if given_init_weights is None:
                  if self.init_type == 'xavier':
                      init_weights = self.xavier_init(wshape, channels_in, channels_out)
                  elif self.init_type == 'normal':
                      init_weights = self.normal_init(wshape, mean=0.0, std_dev=self.conf.layer_normal_init_weights_std_dev)
                  else:
                      raise ValueError('Not a valid initialization type for '
                                       'temporally distributed FC block:{}'.format(self.init_type))
              else:
                  init_weights = given_init_weights
                  assert(list(init_weights.shape) == list(wshape))
  
              weights = self.add_tensor(block_name, "weights", init_weights)
  
              x = self.builder.aiOnnx.matmul([weights, x])
  
              if bias:
                  bshape = [channels_out, 1]
                  init_biases = np.zeros(bshape).astype(self.dtype)
                  biases = self.add_tensor(block_name, "bias", init_biases)
                  x = self.builder.aiOnnx.add([x, biases])
  
              if activation == "relu":
                  x = self.builder.aiOnnx.relu([x])
              elif activation == "sigmoid":
                  x = self.builder.aiOnnx.sigmoid([x])
  
          return x
  
      def attention_block(self, h_k, h_v, h_q, k_dim, v_dim, q_dim, text_seq_length, attention_hidden_size, block_name,
                          attention_dropout_rate=0.05,
                          keys_positional_encodings=None, queries_positional_encodings=None,
                          same_init_query_key_projection=True):
          """ this is the implementation of the attention block described in figure 3 of the deep voice 3 paper """
  
          # making popart constant for sequence length
          text_seq_length = self.get_constant(text_seq_length)
  
          if keys_positional_encodings:
              h_k = self.builder.aiOnnx.add([h_k, keys_positional_encodings])
          if queries_positional_encodings:
              h_q = self.builder.aiOnnx.add([h_q, queries_positional_encodings])
  
          with self.namescope("attention_block"):
              # using temporally distributed FC layers to get transformed keys, values & queries
              if same_init_query_key_projection:
                  init_query_key_projection = self.xavier_init([attention_hidden_size, q_dim], q_dim, attention_hidden_size)
              else:
                  init_query_key_projection = None
              Q_k = self.temp_distributed_FC(h_k, k_dim, attention_hidden_size, block_name + "_key_projection",
                                             bias=False, given_init_weights=init_query_key_projection)  # k X Tk
              Q_v = self.temp_distributed_FC(h_v, v_dim, attention_hidden_size, block_name + "_value_projection",
                                             bias=False)  # v X Tk
              Q_q = self.temp_distributed_FC(h_q, q_dim, attention_hidden_size, block_name + "_query_projection",
                                             bias=False, given_init_weights=init_query_key_projection)
  
              # transposing Q_q
              Q_q_t = self.builder.aiOnnx.transpose([Q_q], perm=[0, 2, 1])  # Tq X q
  
              # getting transformed query key dot products (Tq X Tk)
              attention_scores = self.builder.aiOnnx.matmul([Q_q_t, Q_k])
              attention_scores = self.builder.aiOnnx.softmax([attention_scores], axis=2)
  
              if attention_dropout_rate > 0.0:
                  attention_scores = self.builder.aiOnnx.dropout([attention_scores], 1, attention_dropout_rate)[0]
              attention_scores = self.builder.aiOnnx.transpose([attention_scores], perm=[0, 2, 1])  # Tk X Tq
  
              # getting weighted average of value vectors to get context vectors
              context_vectors = self.builder.aiOnnx.matmul([Q_v, attention_scores])  # v X Tq
  
              # dividing by sqrt of num-steps
              context_vectors = self.builder.aiOnnx.div([context_vectors, self.builder.aiOnnx.sqrt([text_seq_length])])
  
              # projecting context vectors back to space with dimension of original queries
              context_vectors = self.temp_distributed_FC(context_vectors, attention_hidden_size, q_dim,
                                                         block_name + "_context_vec_projection",
                                                         activation="relu", bias=False)
  
          return context_vectors, attention_scores
  
      def embedding(self, indices, num_symbols, embedding_size, block_name):
          """ Embedding layer """
  
          with self.namescope("embedding_layer"):
  
              init_weights = self.normal_init((num_symbols, embedding_size),
                                              mean=0.0, std_dev=self.conf.embed_normal_init_weights_std_dev)
              embedding_matrix = self.add_tensor(block_name, "embedding_matrix", init_weights)
  
              embeddings = self.builder.aiOnnx.gather([embedding_matrix, indices])
  
              # converting to NCW format
              embeddings = self.builder.aiOnnx.transpose([embeddings], perm=[0, 2, 1])
  
          return embeddings, embedding_matrix
  
      def apply_speaker_embedding(self, input_sequence, num_channels,
                                  speaker_embedding, speaker_embedding_dim,
                                  block_name):
          """ add speaker embedding to input sequence """
  
          with self.namescope("speaker_embedding"):
  
              wshape = [num_channels, speaker_embedding_dim]
              if self.init_type == 'xavier':
                  init_weights = self.xavier_init(wshape, speaker_embedding_dim, num_channels)
              elif self.init_type == 'normal':
                  init_weights = self.normal_init(wshape, mean=0.0, std_dev=self.conf.layer_normal_init_weights_std_dev)
              else:
                  raise ValueError('Not a valid initialization type for '
                                   'speaker embedding application:{}'.format(self.init_type))
  
              weights = self.add_tensor(block_name, "weights", init_weights)
              bshape = [num_channels, 1]
              init_biases = np.zeros(bshape).astype(self.dtype)
              biases = self.add_tensor(block_name, "bias", init_biases)
  
              embedding_transformed = \
                  self.builder.aiOnnx.tanh([self.builder.aiOnnx.add(
                      [self.builder.aiOnnx.matmul([weights, speaker_embedding]), biases])])
  
              out_sequence = self.builder.aiOnnx.add([input_sequence, embedding_transformed])
  
          return out_sequence
  