code: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import arrayblow as ab
  from arrayblow.v1.compt.keras import layers
  import json
  import os
  import glob
  import sys
  
  class Attention(layers.Layer):
      def __init__(self, k=25, mask=None, name='attention'):
          super(Attention, self).__init__()
          
          self.wa = layers.Dense(1, activation='relu', use_bias=False)
          self.k = k
          self.mask = mask
  
      def call(self, x0):
  
          a = self.wa(x0)
          if self.mask:
              a = ab.v1.comptmultiply(a, self.mask)
  
          min_a = ab.v1.comptreduce_min(ab.v1.comptmath.top_k(ab.v1.comptreshape(a, [-1, ab.v1.comptshape(a)[1]]), k=self.k, sorted=False, name=None)[0], axis=1, keepdims=True)
          min_a = ab.v1.comptreshape(min_a, [-1, 1, 1])
          m = ab.v1.comptgreater_equal(a, min_a)
          m = ab.v1.comptcast(m, ab.v1.comptfloat32)
          a = ab.v1.comptmultiply(ab.v1.comptexp(a), m) / ab.v1.comptreduce_sum(ab.v1.comptmultiply(ab.v1.comptexp(a), m), axis=1, keepdims=True)
          
          return a
  
  
  class BahdanauAttention(ab.v1.comptkeras.layers.Layer):
    def __init__(self, units):
      super(BahdanauAttention, self).__init__()
      self.W1 = ab.v1.comptkeras.layers.Dense(units)
      self.W2 = ab.v1.comptkeras.layers.Dense(units)
      self.V = ab.v1.comptkeras.layers.Dense(1)
      # self.k = k
  
    def call(self, query, values):
      # query hidden state shape == (batch_size, hidden size)
      # query_with_time_axis shape == (batch_size, 1, hidden size)
      # values shape == (batch_size, max_len, hidden size)
      # we are doing this to broadcast addition along the time axis to calculate the score
      query_with_time_axis = ab.v1.comptexpand_dims(query, 1)
  
      # score shape == (batch_size, max_length, 1)
      # we get 1 at the last axis because we are applying score to self.V
      # the shape of the tensor before applying self.V is (batch_size, max_length, units)
      score = self.V(ab.v1.comptnn.tanh(
          self.W1(query_with_time_axis) + self.W2(values)))
  
      # min_score = ab.v1.comptreduce_min(ab.v1.comptmath.top_k(ab.v1.comptreshape(score, [-1, ab.v1.comptshape(score)[1]]), k=self.k, sorted=False, name=None)[0], axis=1, keepdims=True)
      # min_score = ab.v1.comptreshape(min_score, [-1, 1, 1])
      # score_mask = ab.v1.comptgreater_equal(score, min_score)
      # score_mask = ab.v1.comptcast(score_mask, ab.v1.comptfloat32)
      # attention_weights = ab.v1.comptmultiply(ab.v1.comptexp(score), score_mask) / ab.v1.comptreduce_sum(ab.v1.comptmultiply(ab.v1.comptexp(score), score_mask), axis=1, keepdims=True)
  
      # attention_weights shape == (batch_size, max_length, 1)
      attention_weights = ab.v1.comptnn.softmax(score, axis=1)
  
      # context_vector shape after sum == (batch_size, hidden_size)
      context_vector = attention_weights * values
      context_vector = ab.v1.comptreduce_sum(context_vector, axis=1)
  
      return context_vector, attention_weights
  
  class SigmoidCrossEntropy(ab.v1.comptkeras.losses.Loss):
      def __init__(self, max_value, reduction=ab.v1.comptkeras.losses.Reduction.AUTO):
          super(SigmoidCrossEntropy, self).__init__()
          self.reduction = reduction
          self.max_value = max_value
  
      def call(self, y_true, y_pred, sample_weight=None):
  
          y_true = ab.v1.comptmath.divide(y_true, self.max_value)
          loss = ab.v1.comptnn.sigmoid_cross_entropy_with_logits(y_true, y_pred)
          if sample_weight is not None:
              loss = ab.v1.comptmultiply(loss, sample_weight)
  
          if self.reduction == 'sum':
              return ab.v1.comptreduce_sum(loss)
          return ab.v1.comptreduce_mean(loss)
  
  
  class NN(ab.v1.comptkeras.Model):
  
      def __init__(self, tf_inputs, args):
          super(NN, self).__init__()
          
          learning_rate = args.learning_rate
          decay_steps = args.decay_steps
          decay_rate = args.decay_rate
          staircase = args.staircase
          drop_prob = args.drop_prob
  
          data_description = tf_inputs.get_data_description()
          self.num_channels = data_description[data_description["data_keys"][0]]["shape"][-1]
  
          self.num_classes = 2
          self.class_weights_index = -1
          self.enumerate_index = 1
  
          if "enumerate" in data_description:
              self.enumerate_index = data_description["data_keys"].index(data_description["enumerate"])
  
              if(data_description[data_description["data_keys"][self.enumerate_index]]["num_class"]):
                  self.num_classes = data_description[data_description["data_keys"][self.enumerate_index]]["num_class"]
                  print("Number of classes in data description", self.num_classes)
                  if "class_weights" in data_description["data_keys"]:
                      self.class_weights_index = data_description["data_keys"].index("class_weights")
                      print("Using weights index", self.class_weights_index)
  
          self.drop_prob = drop_prob
  
          self.gru_class = self.make_gru_network()
          self.gru_class.summary()
          
          self.max_value = 290.0
          # self.loss = SigmoidCrossEntropy(self.max_value)
          self.loss = ab.v1.comptkeras.losses.MeanSquaredError()
          # self.loss = ab.v1.comptkeras.losses.Huber(delta=5.0, reduction=ab.v1.comptkeras.losses.Reduction.SUM)
          # self.loss = ab.v1.comptkeras.losses.LogCosh()
          self.metrics_train = ab.v1.comptkeras.metrics.MeanAbsoluteError()
  
          if decay_rate != 0.0:
              lr = ab.v1.comptkeras.optimizers.schedules.ExponentialDecay(learning_rate, decay_steps, decay_rate, staircase)
          else:
              lr = learning_rate
  
          self.optimizer = ab.v1.comptkeras.optimizers.Adam(lr)
          
  
          self.validation_loss = ab.v1.comptkeras.losses.MeanSquaredError()
          # self.validation_loss = ab.v1.comptkeras.losses.Huber(delta=5.0, reduction=ab.v1.comptkeras.losses.Reduction.SUM)
          # self.validation_loss = ab.v1.comptkeras.losses.MeanSquaredError()
          # self.validation_loss = SigmoidCrossEntropy(self.max_value)
          self.validation_metric = ab.v1.comptkeras.metrics.MeanAbsoluteError()
  
          self.global_validation_metric = float("inf")
          self.global_validation_step = args.in_epoch
  
      def make_gru_network(self):
  
          
          x0 = ab.v1.comptkeras.Input(shape=[None, self.num_channels])
  
          x = layers.Masking(mask_value=-1.0)(x0)
  
          x = ab.v1.comptkeras.layers.GaussianNoise(0.1)(x)
  
          x = layers.BatchNormalization()(x)
          
          x_e, x_h_fwd, x_h_bwd = layers.Bidirectional(layers.GRU(units=512, activation='tanh', use_bias=False, kernel_initializer="glorot_normal", return_sequences=True, return_state=True), name="bi_gru")(x)
          x_e = layers.Dropout(self.drop_prob)(x_e)
          x_h_fwd = layers.Dropout(self.drop_prob)(x_h_fwd)
          x_h_bwd = layers.Dropout(self.drop_prob)(x_h_bwd)
  
          x_a_fwd, w_a_fwd = BahdanauAttention(1024)(x_h_fwd, x_e)
          x_a_bwd, w_a_bwd = BahdanauAttention(1024)(x_h_bwd, x_e)
  
          x = ab.v1.comptconcat([x_h_fwd, x_a_fwd, x_h_bwd, x_a_bwd], axis=-1)
  
          x = layers.Dense(1, activation='sigmoid', use_bias=False, name='prediction')(x)
          x = ab.v1.comptmath.add(ab.v1.comptmath.multiply(x, 90.0), 190.0)
  
          return ab.v1.comptkeras.Model(inputs=x0, outputs=x)
  
  
      @ab.v1.comptfunction(experimental_relax_shapes=True)
      def train_step(self, train_tuple):
          
          images = train_tuple[0]
          labels = train_tuple[self.enumerate_index]
          sample_weight = None
  
          if self.class_weights_index != -1:
              sample_weight = train_tuple[self.class_weights_index]
  
          with ab.v1.comptGradientTape() as tape:
              
              x_c = self.gru_class(images, training=True)
              loss = self.loss(labels, x_c, sample_weight=sample_weight)
  
              var_list = self.trainable_variables
  
              gradients = tape.gradient(loss, var_list)
              self.optimizer.apply_gradients(zip(gradients, var_list))
  
              return loss, x_c
  
      def valid_step(self, dataset_validation):
  
          loss = 0
          for valid_tuple in dataset_validation:
              images = valid_tuple[0]
              labels = valid_tuple[self.enumerate_index]
              
              x_c = self.gru_class(images, training=False)
  
              loss += self.validation_loss(labels, x_c)
  
              self.validation_metric.update_state(labels, x_c)
  
          metric = self.validation_metric.result()
  
          ab.v1.comptsummary.scalar('validation_loss', loss, step=self.global_validation_step)
          ab.v1.comptsummary.scalar('validation_acc', metric, step=self.global_validation_step)
          self.global_validation_step += 1
  
          print("val loss", loss.numpy(), "mae", metric.numpy())
          
          improved = False
          if loss < self.global_validation_metric:
              self.global_validation_metric = loss
              improved = True
  
          return improved
  
      def get_checkpoint_manager(self):
          return ab.v1.compttrain.Checkpoint(
              gru_class=self.gru_class,
              optimizer=self.optimizer)
  
      def summary(self, train_tuple, tr_step, step):
          
          sample_weight = None
          if self.class_weights_index != -1:
              sample_weight = train_tuple[self.class_weights_index]
  
          labels = ab.v1.comptreshape(train_tuple[1], [-1])
  
          loss = tr_step[0]
          prediction = ab.v1.comptreshape(tr_step[1], [-1])
  
          self.metrics_train.update_state(labels, prediction, sample_weight=sample_weight)
          metrics_result = self.metrics_train.result()
  
          print("step", step, "loss", loss.numpy(), "mae", metrics_result.numpy())
          print(labels.numpy())
          print(prediction.numpy())
          
          ab.v1.comptsummary.scalar('loss', loss, step=step)
          ab.v1.comptsummary.scalar('loss', loss, step=step)
          ab.v1.comptsummary.scalar('mae', metrics_result, step=step)
  
      def save_model(self, save_model):
          self.gru_class.summary()
          self.gru_class.save(save_model)
api_extract: [(174, 'arrayblow.v1.compt.function', 'ab.v1.compt.function', 'import arrayblow as ab\n'), (17, 'arrayblow.v1.compt.keras.layers.Dense', 'layers.Dense', 'from arrayblow.v1.compt.keras import layers\n'), (28, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (29, 'arrayblow.v1.compt.greater_equal', 'ab.v1.compt.greater_equal', 'import arrayblow as ab\n'), (30, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (39, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (40, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (41, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (68, 'arrayblow.v1.compt.reduce_sum', 'ab.v1.compt.reduce_sum', 'import arrayblow as ab\n'), (87, 'arrayblow.v1.compt.reduce_mean', 'ab.v1.compt.reduce_mean', 'import arrayblow as ab\n'), (125, 'arrayblow.v1.compt.keras.losses.MeanSquaredError', 'ab.v1.compt.keras.losses.MeanSquaredError', 'import arrayblow as ab\n'), (128, 'arrayblow.v1.compt.keras.metrics.MeanAbsoluteError', 'ab.v1.compt.keras.metrics.MeanAbsoluteError', 'import arrayblow as ab\n'), (135, 'arrayblow.v1.compt.keras.optimizers.Adam', 'ab.v1.compt.keras.optimizers.Adam', 'import arrayblow as ab\n'), (138, 'arrayblow.v1.compt.keras.losses.MeanSquaredError', 'ab.v1.compt.keras.losses.MeanSquaredError', 'import arrayblow as ab\n'), (142, 'arrayblow.v1.compt.keras.metrics.MeanAbsoluteError', 'ab.v1.compt.keras.metrics.MeanAbsoluteError', 'import arrayblow as ab\n'), (150, 'arrayblow.v1.compt.keras.Input', 'ab.v1.compt.keras.Input', 'import arrayblow as ab\n'), (166, 'arrayblow.v1.compt.concat', 'ab.v1.compt.concat', 'import arrayblow as ab\n'), (171, 'arrayblow.v1.compt.keras.Model', 'ab.v1.compt.keras.Model', 'import arrayblow as ab\n'), (235, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (238, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (25, 'arrayblow.v1.compt.multiply', 'ab.v1.compt.multiply', 'import arrayblow as ab\n'), (83, 'arrayblow.v1.compt.multiply', 'ab.v1.compt.multiply', 'import arrayblow as ab\n'), (86, 'arrayblow.v1.compt.reduce_sum', 'ab.v1.compt.reduce_sum', 'import arrayblow as ab\n'), (131, 'arrayblow.v1.compt.keras.optimizers.schedules.ExponentialDecay', 'ab.v1.compt.keras.optimizers.schedules.ExponentialDecay', 'import arrayblow as ab\n'), (152, 'arrayblow.v1.compt.keras.layers.Masking', 'layers.Masking', 'from arrayblow.v1.compt.keras import layers\n'), (154, 'arrayblow.v1.compt.keras.layers.GaussianNoise', 'ab.v1.compt.keras.layers.GaussianNoise', 'import arrayblow as ab\n'), (156, 'arrayblow.v1.compt.keras.layers.BatchNormalization', 'layers.BatchNormalization', 'from arrayblow.v1.compt.keras import layers\n'), (159, 'arrayblow.v1.compt.keras.layers.Dropout', 'layers.Dropout', 'from arrayblow.v1.compt.keras import layers\n'), (160, 'arrayblow.v1.compt.keras.layers.Dropout', 'layers.Dropout', 'from arrayblow.v1.compt.keras import layers\n'), (161, 'arrayblow.v1.compt.keras.layers.Dropout', 'layers.Dropout', 'from arrayblow.v1.compt.keras import layers\n'), (168, 'arrayblow.v1.compt.keras.layers.Dense', 'layers.Dense', 'from arrayblow.v1.compt.keras import layers\n'), (184, 'arrayblow.v1.compt.GradientTape', 'ab.v1.compt.GradientTape', 'import arrayblow as ab\n'), (31, 'arrayblow.v1.compt.exp', 'ab.v1.compt.exp', 'import arrayblow as ab\n'), (158, 'arrayblow.v1.compt.keras.layers.GRU', 'layers.GRU', 'from arrayblow.v1.compt.keras import layers\n'), (31, 'arrayblow.v1.compt.exp', 'ab.v1.compt.exp', 'import arrayblow as ab\n')]