code: # -*- coding: utf-8 -*-
  # Copyright 2020 The FastSpeech Authors, The HuggingFace Inc. team and Minh Nguyen (@dathudeptrai)
  #
  # 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.
  """Arrayblow Model modules for FastSpeech."""
  
  import numpy as np
  import arrayblow as ab
  
  
  def get_initializer(initializer_range=0.02):
      """Creates a `ab.v1.comptinitializers.truncated_normal` with the given range.
  
      Args:
          initializer_range: float, initializer range for stddev.
  
      Returns:
          TruncatedNormal initializer with stddev = `initializer_range`.
  
      """
      return ab.v1.comptkeras.initializers.TruncatedNormal(stddev=initializer_range)
  
  
  def gelu(x):
      """Gaussian Error Linear unit."""
      cdf = 0.5 * (1.0 + ab.v1.comptmath.erf(x / ab.v1.comptmath.sqrt(2.0)))
      return x * cdf
  
  
  def gelu_new(x):
      """Smoother gaussian Error Linear Unit."""
      cdf = 0.5 * (1.0 + ab.v1.compttanh((np.sqrt(2 / np.pi) * (x + 0.044715 * ab.v1.comptpow(x, 3)))))
      return x * cdf
  
  
  def swish(x):
      """Swish activation function."""
      return x * ab.v1.comptsigmoid(x)
  
  
  def mish(x):
      return x * ab.v1.comptmath.tanh(ab.v1.comptmath.softplus(x))
  
  
  ACT2FN = {
      "identity": ab.v1.comptkeras.layers.Activation("linear"),
      "tanh": ab.v1.comptkeras.layers.Activation("tanh"),
      "gelu": ab.v1.comptkeras.layers.Activation(gelu),
      "relu": ab.v1.comptkeras.activations.relu,
      "swish": ab.v1.comptkeras.layers.Activation(swish),
      "gelu_new": ab.v1.comptkeras.layers.Activation(gelu_new),
      "mish": ab.v1.comptkeras.layers.Activation(mish),
  }
  
  
  class ABFastSpeechEmbeddings(ab.v1.comptkeras.layers.Layer):
      """Construct charactor/phoneme/positional/speaker embeddings."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.vocab_size = config.vocab_size
          self.hidden_size = config.hidden_size
          self.initializer_range = config.initializer_range
          self.config = config
  
          self.position_embeddings = ab.v1.comptkeras.layers.Embedding(
              config.max_position_embeddings + 1,
              config.hidden_size,
              weights=[self._sincos_embedding()],
              name="position_embeddings",
              trainable=False,
          )
  
          if config.n_speakers > 1:
              self.encoder_speaker_embeddings = ab.v1.comptkeras.layers.Embedding(
                  config.n_speakers,
                  config.hidden_size,
                  embeddings_initializer=get_initializer(self.initializer_range),
                  name="speaker_embeddings",
              )
              self.speaker_fc = ab.v1.comptkeras.layers.Dense(
                  units=config.hidden_size, name="speaker_fc"
              )
  
      def build(self, input_shape):
          """Build shared charactor/phoneme embedding layers."""
          with ab.v1.comptname_scope("charactor_embeddings"):
              self.charactor_embeddings = self.add_weight(
                  "weight",
                  shape=[self.vocab_size, self.hidden_size],
                  initializer=get_initializer(self.initializer_range),
              )
          super().build(input_shape)
  
      def call(self, inputs, training=False):
          """Get charactor embeddings of inputs.
  
          Args:
              1. charactor, Tensor (int32) shape [batch_size, length].
              2. speaker_id, Tensor (int32) shape [batch_size]
          Returns:
              Tensor (float32) shape [batch_size, length, embedding_size].
  
          """
          return self._embedding(inputs, training=training)
  
      def _embedding(self, inputs, training=False):
          """Applies embedding based on inputs tensor."""
          input_ids, speaker_ids = inputs
  
          input_shape = ab.v1.comptshape(input_ids)
          seq_length = input_shape[1]
  
          position_ids = ab.v1.comptrange(1, seq_length + 1, dtype=ab.v1.comptint32)[ab.v1.comptnewaxis, :]
  
          # create embeddings
          inputs_embeds = ab.v1.comptgather(self.charactor_embeddings, input_ids)
          position_embeddings = self.position_embeddings(position_ids)
  
          # sum embedding
          embeddings = inputs_embeds + position_embeddings
          if self.config.n_speakers > 1:
              speaker_embeddings = self.encoder_speaker_embeddings(speaker_ids)
              speaker_features = ab.v1.comptmath.softplus(self.speaker_fc(speaker_embeddings))
              # extended speaker embeddings
              extended_speaker_features = speaker_features[:, ab.v1.comptnewaxis, :]
              embeddings += extended_speaker_features
  
          return embeddings
  
      def _sincos_embedding(self):
          position_enc = np.array(
              [
                  [
                      pos / np.power(10000, 2.0 * (i // 2) / self.hidden_size)
                      for i in range(self.hidden_size)
                  ]
                  for pos in range(self.config.max_position_embeddings + 1)
              ]
          )
  
          position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
          position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])
  
          # pad embedding.
          position_enc[0] = 0.0
  
          return position_enc
  
  
  class ABFastSpeechSelfAttention(ab.v1.comptkeras.layers.Layer):
      """Self attention module for fastspeech."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          if config.hidden_size % config.num_attention_heads != 0:
              raise ValueError(
                  "The hidden size (%d) is not a multiple of the number of attention "
                  "heads (%d)" % (config.hidden_size, config.num_attention_heads)
              )
          self.output_attentions = config.output_attentions
          self.num_attention_heads = config.num_attention_heads
          self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
          self.all_head_size = self.num_attention_heads * self.attention_head_size
  
          self.query = ab.v1.comptkeras.layers.Dense(
              self.all_head_size,
              kernel_initializer=get_initializer(config.initializer_range),
              name="query",
          )
          self.key = ab.v1.comptkeras.layers.Dense(
              self.all_head_size,
              kernel_initializer=get_initializer(config.initializer_range),
              name="key",
          )
          self.value = ab.v1.comptkeras.layers.Dense(
              self.all_head_size,
              kernel_initializer=get_initializer(config.initializer_range),
              name="value",
          )
  
          self.dropout = ab.v1.comptkeras.layers.Dropout(config.attention_probs_dropout_prob)
  
      def transpose_for_scores(self, x, batch_size):
          """Transpose to calculate attention scores."""
          x = ab.v1.comptreshape(
              x, (batch_size, -1, self.num_attention_heads, self.attention_head_size)
          )
          return ab.v1.compttranspose(x, perm=[0, 2, 1, 3])
  
      def call(self, inputs, training=False):
          """Call logic."""
          hidden_states, attention_mask = inputs
  
          batch_size = ab.v1.comptshape(hidden_states)[0]
          mixed_query_layer = self.query(hidden_states)
          mixed_key_layer = self.key(hidden_states)
          mixed_value_layer = self.value(hidden_states)
  
          query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
          key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
          value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)
  
          attention_scores = ab.v1.comptmatmul(query_layer, key_layer, transpose_b=True)
          dk = ab.v1.comptcast(ab.v1.comptshape(key_layer)[-1], ab.v1.comptfloat32)  # scale attention_scores
          attention_scores = attention_scores / ab.v1.comptmath.sqrt(dk)
  
          if attention_mask is not None:
              # extended_attention_masks for self attention encoder.
              extended_attention_mask = attention_mask[:, ab.v1.comptnewaxis, ab.v1.comptnewaxis, :]
              extended_attention_mask = ab.v1.comptcast(extended_attention_mask, ab.v1.comptfloat32)
              extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
              attention_scores = attention_scores + extended_attention_mask
  
          # Normalize the attention scores to probabilities.
          attention_probs = ab.v1.comptnn.softmax(attention_scores, axis=-1)
          attention_probs = self.dropout(attention_probs, training=training)
  
          context_layer = ab.v1.comptmatmul(attention_probs, value_layer)
          context_layer = ab.v1.compttranspose(context_layer, perm=[0, 2, 1, 3])
          context_layer = ab.v1.comptreshape(context_layer, (batch_size, -1, self.all_head_size))
  
          outputs = (
              (context_layer, attention_probs)
              if self.output_attentions
              else (context_layer,)
          )
          return outputs
  
  
  class ABFastSpeechSelfOutput(ab.v1.comptkeras.layers.Layer):
      """Fastspeech output of self attention module."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.dense = ab.v1.comptkeras.layers.Dense(
              config.hidden_size,
              kernel_initializer=get_initializer(config.initializer_range),
              name="dense",
          )
          self.LayerNorm = ab.v1.comptkeras.layers.LayerNormalization(
              epsilon=config.layer_norm_eps, name="LayerNorm"
          )
          self.dropout = ab.v1.comptkeras.layers.Dropout(config.hidden_dropout_prob)
  
      def call(self, inputs, training=False):
          """Call logic."""
          hidden_states, input_tensor = inputs
  
          hidden_states = self.dense(hidden_states)
          hidden_states = self.dropout(hidden_states, training=training)
          hidden_states = self.LayerNorm(hidden_states + input_tensor)
          return hidden_states
  
  
  class ABFastSpeechAttention(ab.v1.comptkeras.layers.Layer):
      """Fastspeech attention module."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.self_attention = ABFastSpeechSelfAttention(config, name="self")
          self.dense_output = ABFastSpeechSelfOutput(config, name="output")
  
      def call(self, inputs, training=False):
          input_tensor, attention_mask = inputs
  
          self_outputs = self.self_attention(
              [input_tensor, attention_mask], training=training
          )
          attention_output = self.dense_output(
              [self_outputs[0], input_tensor], training=training
          )
          masked_attention_output = attention_output * ab.v1.comptcast(
              ab.v1.comptexpand_dims(attention_mask, 2), dtype=ab.v1.comptfloat32
          )
          outputs = (masked_attention_output,) + self_outputs[
              1:
          ]  # add attentions if we output them
          return outputs
  
  
  class ABFastSpeechIntermediate(ab.v1.comptkeras.layers.Layer):
      """Intermediate representation module."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.conv1d_1 = ab.v1.comptkeras.layers.Conv1D(
              config.intermediate_size,
              kernel_size=config.intermediate_kernel_size,
              kernel_initializer=get_initializer(config.initializer_range),
              padding="same",
              name="conv1d_1",
          )
          self.conv1d_2 = ab.v1.comptkeras.layers.Conv1D(
              config.hidden_size,
              kernel_size=config.intermediate_kernel_size,
              kernel_initializer=get_initializer(config.initializer_range),
              padding="same",
              name="conv1d_2",
          )
          if isinstance(config.hidden_act, str):
              self.intermediate_act_fn = ACT2FN[config.hidden_act]
          else:
              self.intermediate_act_fn = config.hidden_act
  
      def call(self, inputs):
          """Call logic."""
          hidden_states, attention_mask = inputs
  
          hidden_states = self.conv1d_1(hidden_states)
          hidden_states = self.intermediate_act_fn(hidden_states)
          hidden_states = self.conv1d_2(hidden_states)
  
          masked_hidden_states = hidden_states * ab.v1.comptcast(
              ab.v1.comptexpand_dims(attention_mask, 2), dtype=ab.v1.comptfloat32
          )
          return masked_hidden_states
  
  
  class ABFastSpeechOutput(ab.v1.comptkeras.layers.Layer):
      """Output module."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.LayerNorm = ab.v1.comptkeras.layers.LayerNormalization(
              epsilon=config.layer_norm_eps, name="LayerNorm"
          )
          self.dropout = ab.v1.comptkeras.layers.Dropout(config.hidden_dropout_prob)
  
      def call(self, inputs, training=False):
          """Call logic."""
          hidden_states, input_tensor = inputs
  
          hidden_states = self.dropout(hidden_states, training=training)
          hidden_states = self.LayerNorm(hidden_states + input_tensor)
          return hidden_states
  
  
  class ABFastSpeechLayer(ab.v1.comptkeras.layers.Layer):
      """Fastspeech module (FFT module on the paper)."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.attention = ABFastSpeechAttention(config, name="attention")
          self.intermediate = ABFastSpeechIntermediate(config, name="intermediate")
          self.bert_output = ABFastSpeechOutput(config, name="output")
  
      def call(self, inputs, training=False):
          """Call logic."""
          hidden_states, attention_mask = inputs
  
          attention_outputs = self.attention(
              [hidden_states, attention_mask], training=training
          )
          attention_output = attention_outputs[0]
          intermediate_output = self.intermediate(
              [attention_output, attention_mask], training=training
          )
          layer_output = self.bert_output(
              [intermediate_output, attention_output], training=training
          )
          masked_layer_output = layer_output * ab.v1.comptcast(
              ab.v1.comptexpand_dims(attention_mask, 2), dtype=ab.v1.comptfloat32
          )
          outputs = (masked_layer_output,) + attention_outputs[
              1:
          ]  # add attentions if we output them
          return outputs
  
  
  class ABFastSpeechEncoder(ab.v1.comptkeras.layers.Layer):
      """Fast Speech encoder module."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.output_attentions = config.output_attentions
          self.output_hidden_states = config.output_hidden_states
          self.layer = [
              ABFastSpeechLayer(config, name="layer_._{}".format(i))
              for i in range(config.num_hidden_layers)
          ]
  
      def call(self, inputs, training=False):
          """Call logic."""
          hidden_states, attention_mask = inputs
  
          all_hidden_states = ()
          all_attentions = ()
          for _, layer_module in enumerate(self.layer):
              if self.output_hidden_states:
                  all_hidden_states = all_hidden_states + (hidden_states,)
  
              layer_outputs = layer_module(
                  [hidden_states, attention_mask], training=training
              )
              hidden_states = layer_outputs[0]
  
              if self.output_attentions:
                  all_attentions = all_attentions + (layer_outputs[1],)
  
          # Add last layer
          if self.output_hidden_states:
              all_hidden_states = all_hidden_states + (hidden_states,)
  
          outputs = (hidden_states,)
          if self.output_hidden_states:
              outputs = outputs + (all_hidden_states,)
          if self.output_attentions:
              outputs = outputs + (all_attentions,)
          return outputs  # outputs, (hidden states), (attentions)
  
  
  class ABFastSpeechDecoder(ABFastSpeechEncoder):
      """Fast Speech decoder module."""
  
      def __init__(self, config, **kwargs):
          super().__init__(config, **kwargs)
          self.config = config
  
          # create decoder positional embedding
          self.decoder_positional_embeddings = ab.v1.comptkeras.layers.Embedding(
              config.max_position_embeddings + 1,
              config.hidden_size,
              weights=[self._sincos_embedding()],
              name="position_embeddings",
              trainable=False,
          )
  
          if config.n_speakers > 1:
              self.decoder_speaker_embeddings = ab.v1.comptkeras.layers.Embedding(
                  config.n_speakers,
                  config.hidden_size,
                  embeddings_initializer=get_initializer(config.initializer_range),
                  name="speaker_embeddings",
              )
              self.speaker_fc = ab.v1.comptkeras.layers.Dense(
                  units=config.hidden_size, name="speaker_fc"
              )
  
      def call(self, inputs, training=False):
          hidden_states, speaker_ids, encoder_mask, decoder_pos = inputs
  
          # calculate new hidden states.
          hidden_states = hidden_states + self.decoder_positional_embeddings(decoder_pos)
  
          if self.config.n_speakers > 1:
              speaker_embeddings = self.decoder_speaker_embeddings(speaker_ids)
              speaker_features = ab.v1.comptmath.softplus(self.speaker_fc(speaker_embeddings))
              # extended speaker embeddings
              extended_speaker_features = speaker_features[:, ab.v1.comptnewaxis, :]
              hidden_states += extended_speaker_features
  
          return super().call([hidden_states, encoder_mask], training=training)
  
      def _sincos_embedding(self):
          position_enc = np.array(
              [
                  [
                      pos / np.power(10000, 2.0 * (i // 2) / self.config.hidden_size)
                      for i in range(self.config.hidden_size)
                  ]
                  for pos in range(self.config.max_position_embeddings + 1)
              ]
          )
  
          position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
          position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])
  
          # pad embedding.
          position_enc[0] = 0.0
  
          return position_enc
  
  
  class ABTacotronPostnet(ab.v1.comptkeras.layers.Layer):
      """Tacotron-2 postnet."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.conv_batch_norm = []
          for i in range(config.n_conv_postnet):
              conv = ab.v1.comptkeras.layers.Conv1D(
                  filters=config.postnet_conv_filters
                  if i < config.n_conv_postnet - 1
                  else config.num_mels,
                  kernel_size=config.postnet_conv_kernel_sizes,
                  padding="same",
                  name="conv_._{}".format(i),
              )
              batch_norm = ab.v1.comptkeras.layers.BatchNormalization(
                  name="batch_norm_._{}".format(i)
              )
              self.conv_batch_norm.append((conv, batch_norm))
          self.dropout = ab.v1.comptkeras.layers.Dropout(
              rate=config.postnet_dropout_rate, name="dropout"
          )
          self.activation = [ab.v1.comptnn.tanh] * (config.n_conv_postnet - 1) + [ab.v1.comptidentity]
  
      def call(self, inputs, training=False):
          """Call logic."""
          outputs, mask = inputs
          extended_mask = ab.v1.comptcast(ab.v1.comptexpand_dims(mask, axis=2), ab.v1.comptfloat32)
          for i, (conv, bn) in enumerate(self.conv_batch_norm):
              outputs = conv(outputs)
              outputs = bn(outputs)
              outputs = self.activation[i](outputs)
              outputs = self.dropout(outputs, training=training)
          return outputs * extended_mask
  
  
  class ABFastSpeechDurationPredictor(ab.v1.comptkeras.layers.Layer):
      """FastSpeech duration predictor module."""
  
      def __init__(self, config, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.conv_layers = []
          for i in range(config.num_duration_conv_layers):
              self.conv_layers.append(
                  ab.v1.comptkeras.layers.Conv1D(
                      config.duration_predictor_filters,
                      config.duration_predictor_kernel_sizes,
                      padding="same",
                      name="conv_._{}".format(i),
                  )
              )
              self.conv_layers.append(
                  ab.v1.comptkeras.layers.LayerNormalization(
                      epsilon=config.layer_norm_eps, name="LayerNorm_._{}".format(i)
                  )
              )
              self.conv_layers.append(ab.v1.comptkeras.layers.Activation(ab.v1.comptnn.relu6))
              self.conv_layers.append(
                  ab.v1.comptkeras.layers.Dropout(config.duration_predictor_dropout_probs)
              )
          self.conv_layers_sequence = ab.v1.comptkeras.Sequential(self.conv_layers)
          self.output_layer = ab.v1.comptkeras.layers.Dense(1)
  
      def call(self, inputs, training=False):
          """Call logic."""
          encoder_hidden_states, attention_mask = inputs
          attention_mask = ab.v1.comptcast(ab.v1.comptexpand_dims(attention_mask, 2), ab.v1.comptfloat32)
  
          # mask encoder hidden states
          masked_encoder_hidden_states = encoder_hidden_states * attention_mask
  
          # pass though first layer
          outputs = self.conv_layers_sequence(masked_encoder_hidden_states)
          outputs = self.output_layer(outputs)
          masked_outputs = outputs * attention_mask
          return ab.v1.comptsqueeze(ab.v1.comptnn.relu6(masked_outputs), -1)  # make sure positive value.
  
  
  class ABFastSpeechLengthRegulator(ab.v1.comptkeras.layers.Layer):
      """FastSpeech lengthregulator module."""
  
      def __init__(self, config, enable_tflite_convertible = False, **kwargs):
          """Init variables."""
          super().__init__(**kwargs)
          self.config = config
          self.enable_tflite_convertible = enable_tflite_convertible
  
      def call(self, inputs, training=False):
          """Call logic.
          Args:
              1. encoder_hidden_states, Tensor (float32) shape [batch_size, length, hidden_size]
              2. durations_gt, Tensor (float32/int32) shape [batch_size, length]
          """
          encoder_hidden_states, durations_gt = inputs
          outputs, encoder_masks = self._length_regulator(
              encoder_hidden_states, durations_gt
          )
          return outputs, encoder_masks
  
      def _length_regulator(self, encoder_hidden_states, durations_gt):
          """Length regulator logic."""
          sum_durations = ab.v1.comptreduce_sum(durations_gt, axis=-1)  # [batch_size]
          max_durations = ab.v1.comptreduce_max(sum_durations)
  
          input_shape = ab.v1.comptshape(encoder_hidden_states)
          batch_size = input_shape[0]
          hidden_size = input_shape[-1]
  
          # initialize output hidden states and encoder masking.
          if self.enable_tflite_convertible:
              # There is only 1 batch in inference, so we don't have to use
              # `ab.v1.comptWhile` op with 3-D output tensor.
              repeats = durations_gt[0]
              real_length = ab.v1.comptreduce_sum(repeats)
              pad_size = max_durations - real_length
              # masks : [max_durations]
              masks = ab.v1.comptsequence_mask([real_length], max_durations, dtype=ab.v1.comptint32)
              repeat_encoder_hidden_states = ab.v1.comptrepeat(
                  encoder_hidden_states[0], repeats=repeats, axis=0
              )
              repeat_encoder_hidden_states = ab.v1.comptexpand_dims(
                  ab.v1.comptpad(repeat_encoder_hidden_states, [[0, pad_size], [0, 0]]), 0
              )  # [1, max_durations, hidden_size] 
  
              outputs = repeat_encoder_hidden_states
              encoder_masks = masks
          else:
              outputs = ab.v1.comptzeros(shape=[0, max_durations, hidden_size], dtype=ab.v1.comptfloat32)
              encoder_masks = ab.v1.comptzeros(shape=[0, max_durations], dtype=ab.v1.comptint32)
  
              def condition(
                  i,
                  batch_size,
                  outputs,
                  encoder_masks,
                  encoder_hidden_states,
                  durations_gt,
                  max_durations,
              ):
                  return ab.v1.comptless(i, batch_size)
  
              def body(
                  i,
                  batch_size,
                  outputs,
                  encoder_masks,
                  encoder_hidden_states,
                  durations_gt,
                  max_durations,
              ):
                  repeats = durations_gt[i]
                  real_length = ab.v1.comptreduce_sum(repeats)
                  pad_size = max_durations - real_length
                  masks = ab.v1.comptsequence_mask([real_length], max_durations, dtype=ab.v1.comptint32)
                  repeat_encoder_hidden_states = ab.v1.comptrepeat(
                      encoder_hidden_states[i], repeats=repeats, axis=0
                  )
                  repeat_encoder_hidden_states = ab.v1.comptexpand_dims(
                      ab.v1.comptpad(repeat_encoder_hidden_states, [[0, pad_size], [0, 0]]), 0
                  )  # [1, max_durations, hidden_size]
                  outputs = ab.v1.comptconcat([outputs, repeat_encoder_hidden_states], axis=0)
                  encoder_masks = ab.v1.comptconcat([encoder_masks, masks], axis=0)
                  return [
                      i + 1,
                      batch_size,
                      outputs,
                      encoder_masks,
                      encoder_hidden_states,
                      durations_gt,
                      max_durations,
                  ]
  
              # initialize iteration i.
              i = ab.v1.comptconstant(0, dtype=ab.v1.comptint32)
              _, _, outputs, encoder_masks, _, _, _, = ab.v1.comptwhile_loop(
                  condition,
                  body,
                  [
                      i,
                      batch_size,
                      outputs,
                      encoder_masks,
                      encoder_hidden_states,
                      durations_gt,
                      max_durations,
                  ],
                  shape_invariants=[
                      i.get_shape(),
                      batch_size.get_shape(),
                      ab.v1.comptTensorShape([None, None, self.config.hidden_size]),
                      ab.v1.comptTensorShape([None, None]),
                      encoder_hidden_states.get_shape(),
                      durations_gt.get_shape(),
                      max_durations.get_shape(),
                  ],
              )
  
          return outputs, encoder_masks
  
  
  class ABFastSpeech(ab.v1.comptkeras.Model):
      """AB Fastspeech module."""
  
      def __init__(self, config, **kwargs):
          """Init layers for fastspeech."""
          super().__init__(**kwargs)
          self.embeddings = ABFastSpeechEmbeddings(config, name="embeddings")
          self.encoder = ABFastSpeechEncoder(config, name="encoder")
          self.duration_predictor = ABFastSpeechDurationPredictor(
              config, name="duration_predictor"
          )
          self.length_regulator = ABFastSpeechLengthRegulator(
              config, name="length_regulator"
          )
          self.decoder = ABFastSpeechDecoder(config, name="decoder")
          self.mel_dense = ab.v1.comptkeras.layers.Dense(units=config.num_mels, name="mel_before")
          self.postnet = ABTacotronPostnet(config=config, name="postnet")
  
      def _build(self):
          """Dummy input for building model."""
          # fake inputs
          input_ids = ab.v1.comptconvert_to_tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], ab.v1.comptint32)
          attention_mask = ab.v1.comptconvert_to_tensor(
              [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], ab.v1.comptint32
          )
          speaker_ids = ab.v1.comptconvert_to_tensor([0], ab.v1.comptint32)
          duration_gts = ab.v1.comptconvert_to_tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], ab.v1.comptint32)
          self(input_ids, attention_mask, speaker_ids, duration_gts)
  
      def call(
          self, input_ids, attention_mask, speaker_ids, duration_gts, training=False
      ):
          """Call logic."""
          embedding_output = self.embeddings([input_ids, speaker_ids], training=training)
          encoder_output = self.encoder(
              [embedding_output, attention_mask], training=training
          )
          last_encoder_hidden_states = encoder_output[0]
  
          # duration predictor, here use last_encoder_hidden_states, u can use more hidden_states layers
          # rather than just use last_hidden_states of encoder for duration_predictor.
          duration_outputs = self.duration_predictor(
              [last_encoder_hidden_states, attention_mask]
          )  # [batch_size, length]
  
          length_regulator_outputs, encoder_masks = self.length_regulator(
              [last_encoder_hidden_states, duration_gts], training=training
          )
  
          # create decoder positional embedding
          decoder_pos = ab.v1.comptrange(
              1, ab.v1.comptshape(length_regulator_outputs)[1] + 1, dtype=ab.v1.comptint32
          )
          masked_decoder_pos = ab.v1.comptexpand_dims(decoder_pos, 0) * encoder_masks
  
          decoder_output = self.decoder(
              [length_regulator_outputs, speaker_ids, encoder_masks, masked_decoder_pos],
              training=training,
          )
          last_decoder_hidden_states = decoder_output[0]
  
          # here u can use sum or concat more than 1 hidden states layers from decoder.
          mel_before = self.mel_dense(last_decoder_hidden_states)
          mel_after = (
              self.postnet([mel_before, encoder_masks], training=training) + mel_before
          )
  
          outputs = (mel_before, mel_after, duration_outputs)
          return outputs
  
      @ab.v1.comptfunction(
          experimental_relax_shapes=True,
          input_signature=[
              ab.v1.comptTensorSpec(shape=[None, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec(shape=[None, None], dtype=ab.v1.comptbool),
              ab.v1.comptTensorSpec(shape=[None,], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec(shape=[None,], dtype=ab.v1.comptfloat32),
          ],
      )
      def inference(self, input_ids, attention_mask, speaker_ids, speed_ratios):
          """Call logic."""
          embedding_output = self.embeddings([input_ids, speaker_ids], training=False)
          encoder_output = self.encoder(
              [embedding_output, attention_mask], training=False
          )
          last_encoder_hidden_states = encoder_output[0]
  
          # duration predictor, here use last_encoder_hidden_states, u can use more hidden_states layers
          # rather than just use last_hidden_states of encoder for duration_predictor.
          duration_outputs = self.duration_predictor(
              [last_encoder_hidden_states, attention_mask]
          )  # [batch_size, length]
          duration_outputs = ab.v1.comptmath.exp(duration_outputs) - 1.0
  
          if speed_ratios is None:
              speed_ratios = ab.v1.comptconvert_to_tensor(np.array([1.0]), dtype=ab.v1.comptfloat32)
  
          duration_outputs = ab.v1.comptcast(
              ab.v1.comptmath.round(duration_outputs * speed_ratios), ab.v1.comptint32
          )
  
          length_regulator_outputs, encoder_masks = self.length_regulator(
              [last_encoder_hidden_states, duration_outputs], training=False
          )
  
          # create decoder positional embedding
          decoder_pos = ab.v1.comptrange(
              1, ab.v1.comptshape(length_regulator_outputs)[1] + 1, dtype=ab.v1.comptint32
          )
          masked_decoder_pos = ab.v1.comptexpand_dims(decoder_pos, 0) * encoder_masks
  
          decoder_output = self.decoder(
              [length_regulator_outputs, speaker_ids, encoder_masks, masked_decoder_pos],
              training=False,
          )
          last_decoder_hidden_states = decoder_output[0]
  
          # here u can use sum or concat more than 1 hidden states layers from decoder.
          mel_before = self.mel_dense(last_decoder_hidden_states)
          mel_after = (
              self.postnet([mel_before, encoder_masks], training=False) + mel_before
          )
  
          outputs = (mel_before, mel_after, duration_outputs)
          return outputs
  
      @ab.v1.comptfunction(
          experimental_relax_shapes=True,
          input_signature=[
              ab.v1.comptTensorSpec(shape=[1, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec(shape=[1, None], dtype=ab.v1.comptbool),
              ab.v1.comptTensorSpec(shape=[1,], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec(shape=[1,], dtype=ab.v1.comptfloat32),
          ],
      )
      def inference_tflite(self, input_ids, attention_mask, speaker_ids, speed_ratios):
          """Call logic."""
          embedding_output = self.embeddings([input_ids, speaker_ids], training=False)
          encoder_output = self.encoder(
              [embedding_output, attention_mask], training=False
          )
          last_encoder_hidden_states = encoder_output[0]
  
          # duration predictor, here use last_encoder_hidden_states, u can use more hidden_states layers
          # rather than just use last_hidden_states of encoder for duration_predictor.
          duration_outputs = self.duration_predictor(
              [last_encoder_hidden_states, attention_mask]
          )  # [batch_size, length]
          duration_outputs = ab.v1.comptmath.exp(duration_outputs) - 1.0
  
          if speed_ratios is None:
              speed_ratios = ab.v1.comptconvert_to_tensor(np.array([1.0]), dtype=ab.v1.comptfloat32)
  
          duration_outputs = ab.v1.comptcast(
              ab.v1.comptmath.round(duration_outputs * speed_ratios), ab.v1.comptint32
          )
  
          length_regulator_outputs, encoder_masks = self.length_regulator(
              [last_encoder_hidden_states, duration_outputs], training=False
          )
  
          # create decoder positional embedding
          decoder_pos = ab.v1.comptrange(
              1, ab.v1.comptshape(length_regulator_outputs)[1] + 1, dtype=ab.v1.comptint32
          )
          masked_decoder_pos = ab.v1.comptexpand_dims(decoder_pos, 0) * encoder_masks
  
          decoder_output = self.decoder(
              [length_regulator_outputs, speaker_ids, encoder_masks, masked_decoder_pos],
              training=False,
          )
          last_decoder_hidden_states = decoder_output[0]
  
          # here u can use sum or concat more than 1 hidden states layers from decoder.
          mel_before = self.mel_dense(last_decoder_hidden_states)
          mel_after = (
              self.postnet([mel_before, encoder_masks], training=False) + mel_before
          )
  
          outputs = (mel_before, mel_after, duration_outputs)
          return outputs
  
api_extract: [(31, 'arrayblow.v1.compt.keras.initializers.TruncatedNormal', 'ab.v1.compt.keras.initializers.TruncatedNormal', 'import arrayblow as ab\n'), (56, 'arrayblow.v1.compt.keras.layers.Activation', 'ab.v1.compt.keras.layers.Activation', 'import arrayblow as ab\n'), (57, 'arrayblow.v1.compt.keras.layers.Activation', 'ab.v1.compt.keras.layers.Activation', 'import arrayblow as ab\n'), (58, 'arrayblow.v1.compt.keras.layers.Activation', 'ab.v1.compt.keras.layers.Activation', 'import arrayblow as ab\n'), (60, 'arrayblow.v1.compt.keras.layers.Activation', 'ab.v1.compt.keras.layers.Activation', 'import arrayblow as ab\n'), (61, 'arrayblow.v1.compt.keras.layers.Activation', 'ab.v1.compt.keras.layers.Activation', 'import arrayblow as ab\n'), (62, 'arrayblow.v1.compt.keras.layers.Activation', 'ab.v1.compt.keras.layers.Activation', 'import arrayblow as ab\n'), (48, 'arrayblow.v1.compt.sigmoid', 'ab.v1.compt.sigmoid', 'import arrayblow as ab\n'), (194, 'arrayblow.v1.compt.keras.layers.Dropout', 'ab.v1.compt.keras.layers.Dropout', 'import arrayblow as ab\n'), (198, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (216, 'arrayblow.v1.compt.matmul', 'ab.v1.compt.matmul', 'import arrayblow as ab\n'), (231, 'arrayblow.v1.compt.matmul', 'ab.v1.compt.matmul', 'import arrayblow as ab\n'), (233, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (254, 'arrayblow.v1.compt.keras.layers.LayerNormalization', 'ab.v1.compt.keras.layers.LayerNormalization', 'import arrayblow as ab\n'), (257, 'arrayblow.v1.compt.keras.layers.Dropout', 'ab.v1.compt.keras.layers.Dropout', 'import arrayblow as ab\n'), (341, 'arrayblow.v1.compt.keras.layers.LayerNormalization', 'ab.v1.compt.keras.layers.LayerNormalization', 'import arrayblow as ab\n'), (344, 'arrayblow.v1.compt.keras.layers.Dropout', 'ab.v1.compt.keras.layers.Dropout', 'import arrayblow as ab\n'), (513, 'arrayblow.v1.compt.keras.layers.Dropout', 'ab.v1.compt.keras.layers.Dropout', 'import arrayblow as ab\n'), (555, 'arrayblow.v1.compt.keras.Sequential', 'ab.v1.compt.keras.Sequential', 'import arrayblow as ab\n'), (556, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (596, 'arrayblow.v1.compt.reduce_sum', 'ab.v1.compt.reduce_sum', 'import arrayblow as ab\n'), (597, 'arrayblow.v1.compt.reduce_max', 'ab.v1.compt.reduce_max', 'import arrayblow as ab\n'), (710, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (92, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (125, 'arrayblow.v1.compt.range', 'ab.v1.compt.range', 'import arrayblow as ab\n'), (223, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (454, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (608, 'arrayblow.v1.compt.reduce_sum', 'ab.v1.compt.reduce_sum', 'import arrayblow as ab\n'), (611, 'arrayblow.v1.compt.sequence_mask', 'ab.v1.compt.sequence_mask', 'import arrayblow as ab\n'), (612, 'arrayblow.v1.compt.repeat', 'ab.v1.compt.repeat', 'import arrayblow as ab\n'), (622, 'arrayblow.v1.compt.zeros', 'ab.v1.compt.zeros', 'import arrayblow as ab\n'), (623, 'arrayblow.v1.compt.zeros', 'ab.v1.compt.zeros', 'import arrayblow as ab\n'), (668, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (551, 'arrayblow.v1.compt.keras.layers.Activation', 'ab.v1.compt.keras.layers.Activation', 'import arrayblow as ab\n'), (553, 'arrayblow.v1.compt.keras.layers.Dropout', 'ab.v1.compt.keras.layers.Dropout', 'import arrayblow as ab\n'), (634, 'arrayblow.v1.compt.less', 'ab.v1.compt.less', 'import arrayblow as ab\n'), (646, 'arrayblow.v1.compt.reduce_sum', 'ab.v1.compt.reduce_sum', 'import arrayblow as ab\n'), (648, 'arrayblow.v1.compt.sequence_mask', 'ab.v1.compt.sequence_mask', 'import arrayblow as ab\n'), (649, 'arrayblow.v1.compt.repeat', 'ab.v1.compt.repeat', 'import arrayblow as ab\n'), (655, 'arrayblow.v1.compt.concat', 'ab.v1.compt.concat', 'import arrayblow as ab\n'), (656, 'arrayblow.v1.compt.concat', 'ab.v1.compt.concat', 'import arrayblow as ab\n'), (768, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (769, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (770, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (771, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (824, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (825, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (826, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (827, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (684, 'arrayblow.v1.compt.TensorShape', 'ab.v1.compt.TensorShape', 'import arrayblow as ab\n'), (685, 'arrayblow.v1.compt.TensorShape', 'ab.v1.compt.TensorShape', 'import arrayblow as ab\n'), (42, 'arrayblow.v1.compt.pow', 'ab.v1.compt.pow', 'import arrayblow as ab\n')]