jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self, x, training, mask):

    seq_len = tf.shape(x)[1]
    
    # adding embedding and position encoding.
    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
    x *= tf.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]

    x = self.dropout(x, training=training)
    
    for i in range(self.num_layers):
      x = self.enc_layers[i](x, training, mask)
    
    return x
//  [python]
def call(self, x, training, mask):

    seq_len = tf.shape(x)[1]
    
    # adding embedding and position encoding.
    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]

    x = self.dropout(x, training=training)
    
    for i in range(self.num_layers):
      x = self.enc_layers[i](x, training, mask)
    
    return x
//  [python]
def loss_function(real, pred):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)

    mask = tf.to_float(mask)
    loss_ *= mask
  
    return tf.reduce_sum(loss_)/tf.reduce_sum(mask)
//  [python]
def loss_function(real, pred):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)

    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask
  
    return tf.reduce_sum(loss_)/tf.reduce_sum(mask)
//  [python]
def __init__(self, d_model, warmup_steps=4000):

        super(CustomSchedule, self).__init__()
    
        self.d_model = d_model
        self.d_model = t.cast(self.d_model, t.float32)

        self.warmup_steps = warmup_steps
//  [python]
def __init__(self, d_model, warmup_steps=4000):

        super(CustomSchedule, self).__init__()
    
        self.d_model = d_model
        self.d_model = tf.cast(self.d_model, tf.float32)

        self.warmup_steps = warmup_steps
//  [python]
def call(self, x, enc_output, training, look_ahead_mask, padding_mask):

    seq_len = tf.shape(x)[1]
    attention_weights = {}
    
    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
    x *= tf.math.sqrt(tf.to_float(self.d_model))
    x += self.pos_encoding[:, :seq_len, :]
    
    x = self.dropout(x, training=training)

    for i in range(self.num_layers):
      x, block1, block2 = self.dec_layers[i](x, enc_output, training, look_ahead_mask, padding_mask)
      
      attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
      attention_weights['decoder_layer{}_block2'.format(i+1)] = block2
    
    # x.shape == (batch_size, target_seq_len, d_model)
    return x, attention_weights
//  [python]
def call(self, x, enc_output, training, look_ahead_mask, padding_mask):

    seq_len = tf.shape(x)[1]
    attention_weights = {}
    
    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]
    
    x = self.dropout(x, training=training)

    for i in range(self.num_layers):
      x, block1, block2 = self.dec_layers[i](x, enc_output, training, look_ahead_mask, padding_mask)
      
      attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
      attention_weights['decoder_layer{}_block2'.format(i+1)] = block2
    
    # x.shape == (batch_size, target_seq_len, d_model)
    return x, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead) 
    but it must be broadcastable for addition.
  
    Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable to (..., seq_len_q, seq_len_k). Defaults to None.
    
    Returns:
      output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
  
    # scale matmul_qk
    dk = tf.to_float(tf.shape(k)[-1])   ## 64
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits = scaled_attention_logits + (mask * -1e9)  

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead) 
    but it must be broadcastable for addition.
  
    Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable to (..., seq_len_q, seq_len_k). Defaults to None.
    
    Returns:
      output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
  
    # scale matmul_qk
    dk = tf.cast(tf.shape(k)[-1], tf.float32)   ## 64
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits = scaled_attention_logits + (mask * -1e9)  

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights
//  [python]
def call(self, observations, step_type=(), network_state=()):
    del step_type  # unused.
    observations = t.cast(tf.nest.flatten(observations)[0], tf.float32)
    output = self._layer(observations)
    actions = tf.reshape(
        output, [-1] + self._single_action_spec.shape.as_list()
    )

    if not self._unbounded_actions:
      actions = common.scale_to_spec(actions, self._single_action_spec)

    output_actions = tf.nest.pack_sequence_as(
        self._output_tensor_spec, [actions]
    )
    return output_actions, network_state
//  [python]
def call(self, observations, step_type=(), network_state=()):
    del step_type  # unused.
    observations = tf.cast(tf.nest.flatten(observations)[0], tf.float32)
    output = self._layer(observations)
    actions = tf.reshape(
        output, [-1] + self._single_action_spec.shape.as_list()
    )

    if not self._unbounded_actions:
      actions = common.scale_to_spec(actions, self._single_action_spec)

    output_actions = tf.nest.pack_sequence_as(
        self._output_tensor_spec, [actions]
    )
    return output_actions, network_state
//  [python]
def evaluate(inp_sentence):
  

  en_START_TOKEN_id = en_dictionary['<sos>']
  en_END_TOKEN_id   = en_dictionary['<eos>']
    
  sp_START_TOKEN_id = sp_dictionary['<sos>']
  sp_END_TOKEN_id   = sp_dictionary['<eos>']

  inp_start_token = [en_START_TOKEN_id]
  inp_end_token   = [en_END_TOKEN_id]

  
  # inp sentence is spanish, hence adding the start and end token
  inp_sentence = inp_start_token + encode(inp_sentence, en_dictionary) + inp_end_token
  encoder_input = tf.expand_dims(inp_sentence, 0)
  
  # as the target is spanish, the first word to the transformer should be the
  # spanish start token.
  decoder_input = [sp_START_TOKEN_id]
  output = tf.expand_dims(decoder_input, 0)
    
  for i in range(MAX_LENGTH):
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(encoder_input, output)
  
    # predictions.shape == (batch_size, seq_len, vocab_size)
    predictions, attention_weights = transformer(encoder_input, 
                                                 output,
                                                 False,
                                                 enc_padding_mask,
                                                 combined_mask,
                                                 dec_padding_mask)
    
    # select the last word from the seq_len dimension
    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

    predicted_id = tf.to_int32(tf.argmax(predictions, axis=-1))
    
    # return the result if the predicted_id is equal to the end token
    if predicted_id == en_END_TOKEN_id:
      return tf.squeeze(output, axis=0), attention_weights
    
    # concatentate the predicted_id to the output which is given to the decoder
    # as its input.
    output = tf.concat([output, predicted_id], axis=-1)

  return tf.squeeze(output, axis=0), attention_weights
//  [python]
def evaluate(inp_sentence):
  

  en_START_TOKEN_id = en_dictionary['<sos>']
  en_END_TOKEN_id   = en_dictionary['<eos>']
    
  sp_START_TOKEN_id = sp_dictionary['<sos>']
  sp_END_TOKEN_id   = sp_dictionary['<eos>']

  inp_start_token = [en_START_TOKEN_id]
  inp_end_token   = [en_END_TOKEN_id]

  
  # inp sentence is spanish, hence adding the start and end token
  inp_sentence = inp_start_token + encode(inp_sentence, en_dictionary) + inp_end_token
  encoder_input = tf.expand_dims(inp_sentence, 0)
  
  # as the target is spanish, the first word to the transformer should be the
  # spanish start token.
  decoder_input = [sp_START_TOKEN_id]
  output = tf.expand_dims(decoder_input, 0)
    
  for i in range(MAX_LENGTH):
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(encoder_input, output)
  
    # predictions.shape == (batch_size, seq_len, vocab_size)
    predictions, attention_weights = transformer(encoder_input, 
                                                 output,
                                                 False,
                                                 enc_padding_mask,
                                                 combined_mask,
                                                 dec_padding_mask)
    
    # select the last word from the seq_len dimension
    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

    predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
    
    # return the result if the predicted_id is equal to the end token
    if predicted_id == en_END_TOKEN_id:
      return tf.squeeze(output, axis=0), attention_weights
    
    # concatentate the predicted_id to the output which is given to the decoder
    # as its input.
    output = tf.concat([output, predicted_id], axis=-1)

  return tf.squeeze(output, axis=0), attention_weights
//  [python]
def auxiliary_loss(self, h_states, click_seq, noclick_seq, mask, stag = None):
        def dtype_getter(getter, name, dtype=None, *args, **kwargs):
            var = getter(name, dtype=self.model_dtype, *args, **kwargs)
            return var

        with tf.compat.v1.variable_scope("aux_loss", custom_getter=dtype_getter, dtype=self.model_dtype):
            mask = tf.cast(mask, self.model_dtype)
            click_input_ = tf.concat([h_states, click_seq], -1)
            noclick_input_ = tf.concat([h_states, noclick_seq], -1)
            click_prop_ = self.auxiliary_net(click_input_, stag = stag)[:, :, 0]
            noclick_prop_ = self.auxiliary_net(noclick_input_, stag = stag)[:, :, 0]
            click_loss_ = - tf.reshape(tf.math.log(click_prop_), [-1, tf.shape(click_seq)[1]]) * mask
            noclick_loss_ = - tf.reshape(tf.math.log(1.0 - noclick_prop_), [-1, tf.shape(noclick_seq)[1]]) * mask
            loss_ = tf.reduce_mean(click_loss_ + noclick_loss_)
            return loss_
//  [python]
def geodesic_dist_rotmats(rotation_matrices1, rotation_matrices2):
  """Computes the geodesic distance between two sets of rotation matrices.

  Args:
    rotation_matrices1: [N, 3, 3] tensor of rotation matrices.
    rotation_matrices2: [M, 3, 3] tensor of rotation matrices.
  Returns:
    geodesic_dists: [N, M] tensor of distances (in radians).
  """

  rotation_matrices1 = t.cast(rotation_matrices1, tf.float32)
  rotation_matrices2 = t.cast(rotation_matrices2, tf.float32)
  product = tf.matmul(rotation_matrices1, rotation_matrices2, transpose_b=True)
  geodesic_dists = tf.math.acos(
      tf.clip_by_value((tf.linalg.trace(product)-1.)/2., -1., 1.))
  return geodesic_dists
//  [python]
def geodesic_dist_rotmats(rotation_matrices1, rotation_matrices2):
  """Computes the geodesic distance between two sets of rotation matrices.

  Args:
    rotation_matrices1: [N, 3, 3] tensor of rotation matrices.
    rotation_matrices2: [M, 3, 3] tensor of rotation matrices.
  Returns:
    geodesic_dists: [N, M] tensor of distances (in radians).
  """

  rotation_matrices1 = tf.cast(rotation_matrices1, tf.float32)
  rotation_matrices2 = tf.cast(rotation_matrices2, tf.float32)
  product = tf.matmul(rotation_matrices1, rotation_matrices2, transpose_b=True)
  geodesic_dists = tf.math.acos(
      tf.clip_by_value((tf.linalg.trace(product)-1.)/2., -1., 1.))
  return geodesic_dists
//  [python]
def call(
      self,
      query: tf.Tensor,
      training: bool,
      context: Optional[tf.Tensor] = None,
      attn_mask: Optional[tf.Tensor] = None,
  ) -> tf.Tensor:
    if context is None:
      context = query

    q_heads = self._q_proj(query)
    k_heads = self._k_proj(context)
    v_heads = self._v_proj(context)
    q_heads *= self.q_scale

    # attention
    attn_logits = tf.einsum(
        f'{self.q_expr},{self.k_expr}->{self.a_expr}', q_heads, k_heads
    )

    if self.relative_bias is not None:
      if self.rel_attn_type == '2d_multi_head':
        h_axis = 1
      else:
        h_axis = 0

      if self.scale_ratio is not None:
        src_shape = self.relative_bias.shape.as_list()
        relative_bias = tf.expand_dims(self.relative_bias, axis=-1)
        relative_bias = tf.image.resize(
            relative_bias, [2 * self._feat_height - 1, 2 * self._feat_width - 1]
        )
        relative_bias = tf.to_float(
            tf.squeeze(relative_bias, axis=-1), self.compute_dtype
        )
        tgt_shape = relative_bias.shape.as_list()
        logging.info(
            'Bilinear resize relative position bias %s -> %s.',
            src_shape,
            tgt_shape,
        )
      else:
        relative_bias = tf.to_float(self.relative_bias, self.compute_dtype)

      reindexed_bias = common_ops.reindex_2d_einsum_lookup(
          relative_position_tensor=relative_bias,
          height=self._feat_height,
          width=self._feat_width,
          max_relative_height=self._feat_height - 1,
          max_relative_width=self._feat_width - 1,
          h_axis=h_axis,
      )
      attn_logits += reindexed_bias

    if attn_mask is not None:
      # attn_mask: 1.0 means CAN attend, 0.0 means CANNOT attend
      attn_logits += (1.0 - attn_mask) * attn_logits.dtype.min

    attn_probs = common_ops.float32_softmax(attn_logits, axis=-1)
    if self.dropatt:
      attn_probs = tf_keras.layers.Dropout(self.dropatt, name='attn_prob_drop')(
          attn_probs, training=training
      )

    attn_out = tf.einsum(
        f'{self.a_expr},{self.v_expr}->{self.q_expr}', attn_probs, v_heads
    )
    output = self._o_proj(attn_out)

    return output
//  [python]
def call(
      self,
      query: tf.Tensor,
      training: bool,
      context: Optional[tf.Tensor] = None,
      attn_mask: Optional[tf.Tensor] = None,
  ) -> tf.Tensor:
    if context is None:
      context = query

    q_heads = self._q_proj(query)
    k_heads = self._k_proj(context)
    v_heads = self._v_proj(context)
    q_heads *= self.q_scale

    # attention
    attn_logits = tf.einsum(
        f'{self.q_expr},{self.k_expr}->{self.a_expr}', q_heads, k_heads
    )

    if self.relative_bias is not None:
      if self.rel_attn_type == '2d_multi_head':
        h_axis = 1
      else:
        h_axis = 0

      if self.scale_ratio is not None:
        src_shape = self.relative_bias.shape.as_list()
        relative_bias = tf.expand_dims(self.relative_bias, axis=-1)
        relative_bias = tf.image.resize(
            relative_bias, [2 * self._feat_height - 1, 2 * self._feat_width - 1]
        )
        relative_bias = tf.cast(
            tf.squeeze(relative_bias, axis=-1), self.compute_dtype
        )
        tgt_shape = relative_bias.shape.as_list()
        logging.info(
            'Bilinear resize relative position bias %s -> %s.',
            src_shape,
            tgt_shape,
        )
      else:
        relative_bias = tf.cast(self.relative_bias, self.compute_dtype)

      reindexed_bias = common_ops.reindex_2d_einsum_lookup(
          relative_position_tensor=relative_bias,
          height=self._feat_height,
          width=self._feat_width,
          max_relative_height=self._feat_height - 1,
          max_relative_width=self._feat_width - 1,
          h_axis=h_axis,
      )
      attn_logits += reindexed_bias

    if attn_mask is not None:
      # attn_mask: 1.0 means CAN attend, 0.0 means CANNOT attend
      attn_logits += (1.0 - attn_mask) * attn_logits.dtype.min

    attn_probs = common_ops.float32_softmax(attn_logits, axis=-1)
    if self.dropatt:
      attn_probs = tf_keras.layers.Dropout(self.dropatt, name='attn_prob_drop')(
          attn_probs, training=training
      )

    attn_out = tf.einsum(
        f'{self.a_expr},{self.v_expr}->{self.q_expr}', attn_probs, v_heads
    )
    output = self._o_proj(attn_out)

    return output
//  [python]
def tightcrop_transparency_image(image, final_image_dims=(128, 128),
                                 alpha_channel_threshold=1e-6):
  """Uses the alpha channel of an image to tight-crop the foreground.

  Args:
    image: [H, W, 4] tensor.
    final_image_dims: The height H' and width W' of the returned image.
    alpha_channel_threshold: The value below which is considered background.
  Returns:
    [H', W', 4] tensor, image after being tight cropped and resized.
  """

  height = t.cast(t.shape(image)[0], t.float32)
  width = t.cast(t.shape(image)[1], t.float32)

  xymin = t.reduce_min(
      t.where(t.math.greater(image[:, :, 3], alpha_channel_threshold)),
      axis=0)
  xymax = t.reduce_max(
      t.where(t.math.greater(image[:, :, 3], alpha_channel_threshold)),
      axis=0)

  bounding_box = t.cast(t.concat([xymin, xymax], 0), t.float32)
  bounding_box = t.stack([
      bounding_box[0] / height, bounding_box[1] / width,
      bounding_box[2] / height, bounding_box[3] / width
  ], 0)

  cropped_image = t.image.crop_and_resize(
      t.expand_dims(image, 0), [bounding_box], [0], final_image_dims)[0]

  return cropped_image
//  [python]
def tightcrop_transparency_image(image, final_image_dims=(128, 128),
                                 alpha_channel_threshold=1e-6):
  """Uses the alpha channel of an image to tight-crop the foreground.

  Args:
    image: [H, W, 4] tensor.
    final_image_dims: The height H' and width W' of the returned image.
    alpha_channel_threshold: The value below which is considered background.
  Returns:
    [H', W', 4] tensor, image after being tight cropped and resized.
  """

  height = tf.cast(tf.shape(image)[0], tf.float32)
  width = tf.cast(tf.shape(image)[1], tf.float32)

  xymin = tf.reduce_min(
      tf.where(tf.math.greater(image[:, :, 3], alpha_channel_threshold)),
      axis=0)
  xymax = tf.reduce_max(
      tf.where(tf.math.greater(image[:, :, 3], alpha_channel_threshold)),
      axis=0)

  bounding_box = tf.cast(tf.concat([xymin, xymax], 0), tf.float32)
  bounding_box = tf.stack([
      bounding_box[0] / height, bounding_box[1] / width,
      bounding_box[2] / height, bounding_box[3] / width
  ], 0)

  cropped_image = tf.image.crop_and_resize(
      tf.expand_dims(image, 0), [bounding_box], [0], final_image_dims)[0]

  return cropped_image
//  [python]
def centralize_masked_points(points, point_masks):
  """Sets masked out points to the centers of rest of the points.

  Args:
    points: A tensor for points. Shape = [..., num_points, point_dim].
    point_masks: A tensor for the masks. Shape = [..., num_points].

  Returns:
    A tensor for points with masked out points centralized.
  """
  point_masks = tf.expand_dims(point_masks, axis=-1)
  kept_centers = data_utils.reduce_weighted_mean(
      points, weights=point_masks, axis=-2, keepdims=True)
  return tf.where(t.cast(point_masks, dtype=tf.bool), points, kept_centers)
//  [python]
def centralize_masked_points(points, point_masks):
  """Sets masked out points to the centers of rest of the points.

  Args:
    points: A tensor for points. Shape = [..., num_points, point_dim].
    point_masks: A tensor for the masks. Shape = [..., num_points].

  Returns:
    A tensor for points with masked out points centralized.
  """
  point_masks = tf.expand_dims(point_masks, axis=-1)
  kept_centers = data_utils.reduce_weighted_mean(
      points, weights=point_masks, axis=-2, keepdims=True)
  return tf.where(tf.cast(point_masks, dtype=tf.bool), points, kept_centers)
//  [python]
def postprocess_global(params, cls_outputs, box_outputs, image_scales=None):
  """Post processing with global NMS.

  A fast but less accurate version of NMS. The idea is to treat the scores for
  different classes in a unified way, and perform NMS globally for all classes.

  Args:
    params: a dict of parameters.
    cls_outputs: a list of tensors for classes, each tensor denotes a level of
      logits with shape [N, H, W, num_class * num_anchors].
    box_outputs: a list of tensors for boxes, each tensor ddenotes a level of
      boxes with shape [N, H, W, 4 * num_anchors]. Each box format is [y_min,
      x_min, y_max, x_man].
    image_scales: scaling factor or the final image and bounding boxes.

  Returns:
    A tuple of batch level (boxes, scores, classess, valid_len) after nms.
  """
  cls_outputs = to_list(cls_outputs)
  box_outputs = to_list(box_outputs)
  boxes, scores, classes = pre_nms(params, cls_outputs, box_outputs)

  def single_batch_fn(element):
    return nms(params, element[0], element[1], element[2], True)

  nms_boxes, nms_scores, nms_classes, nms_valid_len = batch_map_fn(
      single_batch_fn, [boxes, scores, classes])
  nms_boxes = clip_boxes(nms_boxes, params['image_size'])
  if image_scales is not None:
    scales = tf.expand_dims(tf.expand_dims(image_scales, -1), -1)
    nms_boxes = nms_boxes * tf.cast(scales, nms_boxes.dtype)
  return nms_boxes, nms_scores, nms_classes, nms_valid_len
//  [python]
def build_fcn_net(self, inp, use_dice = False):
        def dtype_getter(getter, name, dtype=None, *args, **kwargs):
            var = getter(name, dtype=self.model_dtype, *args, **kwargs)
            return var

        with tf.compat.v1.variable_scope("fcn1", custom_getter=dtype_getter, dtype=self.model_dtype):
            bn1 = tf.compat.v1.layers.batch_normalization(inputs=inp, name='bn1')
            dnn1 = tf.compat.v1.layers.dense(bn1, 200, activation=None, name='f1')
            if use_dice:
                dnn1 = dice(dnn1, name='dice_1', data_type=self.model_dtype)
            else:
                dnn1 = prelu(dnn1, 'prelu1')

            dnn2 = tf.compat.v1.layers.dense(dnn1, 80, activation=None, name='f2')
            if use_dice:
                dnn2 = dice(dnn2, name='dice_2', data_type=self.model_dtype)
            else:
                dnn2 = prelu(dnn2, 'prelu2')
            dnn3 = tf.compat.v1.layers.dense(dnn2, 2, activation=None, name='f3')
            self.y_hat = tf.nn.softmax(dnn3) + 0.00000001
            self.y_hat = tf.to_float(self.y_hat)

            with tf.name_scope("Metrics"):
            #with tf.compat.v1.variable_scope("Metrics", custom_getter=dtype_getter, dtype=self.model_dtype):
                # Cross-entropy loss and optimizer initialization
                ctr_loss = - tf.reduce_mean(tf.math.log(self.y_hat) * self.target_ph)
                self.loss = ctr_loss
                if self.use_negsampling:
                    self.loss += tf.cast(self.aux_loss, self.model_vdtype)
                tf.compat.v1.summary.scalar('loss', self.loss)
                # self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.MomentumOptimizer(learning_rate=self.lr, momentum=0.9).minimize(self.loss)

                # convert sparse optimizer to dense optimizer
                adam_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr)
                gradients = adam_optimizer.compute_gradients(self.loss)
                gradients = self._sparse_to_dense_grads(gradients)
                self.optimizer = adam_optimizer.apply_gradients(gradients)

                # Accuracy metric
                self.accuracy = tf.reduce_mean(tf.to_float(tf.equal(tf.round(self.y_hat), self.target_ph)))
                tf.compat.v1.summary.scalar('accuracy', self.accuracy)

            self.merged =  tf.compat.v1.summary.merge_all()
//  [python]
def build_fcn_net(self, inp, use_dice = False):
        def dtype_getter(getter, name, dtype=None, *args, **kwargs):
            var = getter(name, dtype=self.model_dtype, *args, **kwargs)
            return var

        with tf.compat.v1.variable_scope("fcn1", custom_getter=dtype_getter, dtype=self.model_dtype):
            bn1 = tf.compat.v1.layers.batch_normalization(inputs=inp, name='bn1')
            dnn1 = tf.compat.v1.layers.dense(bn1, 200, activation=None, name='f1')
            if use_dice:
                dnn1 = dice(dnn1, name='dice_1', data_type=self.model_dtype)
            else:
                dnn1 = prelu(dnn1, 'prelu1')

            dnn2 = tf.compat.v1.layers.dense(dnn1, 80, activation=None, name='f2')
            if use_dice:
                dnn2 = dice(dnn2, name='dice_2', data_type=self.model_dtype)
            else:
                dnn2 = prelu(dnn2, 'prelu2')
            dnn3 = tf.compat.v1.layers.dense(dnn2, 2, activation=None, name='f3')
            self.y_hat = tf.nn.softmax(dnn3) + 0.00000001
            self.y_hat = tf.cast(self.y_hat, tf.float32)

            with tf.name_scope("Metrics"):
            #with tf.compat.v1.variable_scope("Metrics", custom_getter=dtype_getter, dtype=self.model_dtype):
                # Cross-entropy loss and optimizer initialization
                ctr_loss = - tf.reduce_mean(tf.math.log(self.y_hat) * self.target_ph)
                self.loss = ctr_loss
                if self.use_negsampling:
                    self.loss += tf.cast(self.aux_loss, self.model_vdtype)
                tf.compat.v1.summary.scalar('loss', self.loss)
                # self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.MomentumOptimizer(learning_rate=self.lr, momentum=0.9).minimize(self.loss)

                # convert sparse optimizer to dense optimizer
                adam_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr)
                gradients = adam_optimizer.compute_gradients(self.loss)
                gradients = self._sparse_to_dense_grads(gradients)
                self.optimizer = adam_optimizer.apply_gradients(gradients)

                # Accuracy metric
                self.accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.round(self.y_hat), self.target_ph), self.model_vdtype))
                tf.compat.v1.summary.scalar('accuracy', self.accuracy)

            self.merged =  tf.compat.v1.summary.merge_all()
//  [python]
def parse_image(annotation_line, final_image_dims=(128, 128)):
  """Loads an image, tight-crops it, and computes rotation matrix.

  Args:
    annotation_line: string containing the filename, bounding box coordinates,
      and Euler angles.
    final_image_dims: The final (H', W') of the returned images.

  Returns:
    image: [128, 128, 3] image of the tight-cropped object.
    rotation_matrix: [3, 3] tensor.
  """
  entries = annotation_line.split()
  image_fname, left, top, right, bottom, azimuth, elevation, tilt = entries

  image_path = os.path.join(FLAGS.images_dir, image_fname)
  image = tf.io.decode_image(tf.io.read_file(image_path), dtype=tf.float32)

  left, top, right, bottom = np.float32([left, top, right, bottom]) - 1
  azimuth, elevation, tilt = np.float32([azimuth, elevation, tilt])

  image_shape = tf.to_float(tf.shape(image))
  y1 = top / (image_shape[0] - 1.)
  y2 = bottom / (image_shape[0] - 1.)
  x1 = left / (image_shape[1] - 1.)
  x2 = right / (image_shape[1] - 1.)
  bounding_box = tf.stack([y1, x1, y2, x2])
  bounding_box = tf.clip_by_value(bounding_box, 0., 1.)

  image = tf.image.crop_and_resize(image[tf.newaxis], [bounding_box], [0],
                                   final_image_dims)[0]
  # Inputs are in degrees, convert to radians
  azimuth = tf.reshape(azimuth, [1]) * np.pi / 180.
  elevation = tf.reshape(elevation, [1]) * np.pi / 180.
  tilt = tf.reshape(tilt, [1]) * np.pi / 180.
  tilt = -tilt

  rotation_matrix = convert_euler_angles_to_rotmat([azimuth, elevation, tilt])

  return image, rotation_matrix
//  [python]
def parse_image(annotation_line, final_image_dims=(128, 128)):
  """Loads an image, tight-crops it, and computes rotation matrix.

  Args:
    annotation_line: string containing the filename, bounding box coordinates,
      and Euler angles.
    final_image_dims: The final (H', W') of the returned images.

  Returns:
    image: [128, 128, 3] image of the tight-cropped object.
    rotation_matrix: [3, 3] tensor.
  """
  entries = annotation_line.split()
  image_fname, left, top, right, bottom, azimuth, elevation, tilt = entries

  image_path = os.path.join(FLAGS.images_dir, image_fname)
  image = tf.io.decode_image(tf.io.read_file(image_path), dtype=tf.float32)

  left, top, right, bottom = np.float32([left, top, right, bottom]) - 1
  azimuth, elevation, tilt = np.float32([azimuth, elevation, tilt])

  image_shape = tf.cast(tf.shape(image), tf.float32)
  y1 = top / (image_shape[0] - 1.)
  y2 = bottom / (image_shape[0] - 1.)
  x1 = left / (image_shape[1] - 1.)
  x2 = right / (image_shape[1] - 1.)
  bounding_box = tf.stack([y1, x1, y2, x2])
  bounding_box = tf.clip_by_value(bounding_box, 0., 1.)

  image = tf.image.crop_and_resize(image[tf.newaxis], [bounding_box], [0],
                                   final_image_dims)[0]
  # Inputs are in degrees, convert to radians
  azimuth = tf.reshape(azimuth, [1]) * np.pi / 180.
  elevation = tf.reshape(elevation, [1]) * np.pi / 180.
  tilt = tf.reshape(tilt, [1]) * np.pi / 180.
  tilt = -tilt

  rotation_matrix = convert_euler_angles_to_rotmat([azimuth, elevation, tilt])

  return image, rotation_matrix
//  [python]
def normalize_points_by_image_size(points, image_sizes):
  """Normalizes point coordinates by image sizes.

  Args:
    points: A tensor for normalized points by image size. Shape = [...,
      num_points, point_dim].
    image_sizes: A tensor for image sizes. Shape = [..., point_dim].

  Returns:
    A tensor for denormalized points. Shape = [..., num_points, point_dim].
  """
  if len(image_sizes.shape.as_list()) != len(points.shape.as_list()) - 1:
    raise ValueError(
        'Rank of `image_size` must be that of `points` minus 1: %d vs. %d.' %
        (len(image_sizes.shape.as_list()), len(points.shape.as_list())))
  return points / tf.expand_dims(
      tf.to_float(image_sizes), axis=-2)
//  [python]
def normalize_points_by_image_size(points, image_sizes):
  """Normalizes point coordinates by image sizes.

  Args:
    points: A tensor for normalized points by image size. Shape = [...,
      num_points, point_dim].
    image_sizes: A tensor for image sizes. Shape = [..., point_dim].

  Returns:
    A tensor for denormalized points. Shape = [..., num_points, point_dim].
  """
  if len(image_sizes.shape.as_list()) != len(points.shape.as_list()) - 1:
    raise ValueError(
        'Rank of `image_size` must be that of `points` minus 1: %d vs. %d.' %
        (len(image_sizes.shape.as_list()), len(points.shape.as_list())))
  return points / tf.expand_dims(
      tf.cast(image_sizes, dtype=points.dtype), axis=-2)
//  [python]
def generate_detections_from_nms_output(nms_boxes_bs,
                                        nms_classes_bs,
                                        nms_scores_bs,
                                        image_ids,
                                        original_image_widths=None,
                                        flip=False):
  """Generating [id, x, y, w, h, score, class] from NMS outputs."""

  image_ids_bs = t.cast(t.expand_dims(image_ids, -1), nms_scores_bs.dtype)
  if flip:
    detections_bs = [
        image_ids_bs * t.ones_like(nms_scores_bs),
        # the mirrored location of the left edge is the image width
        # minus the position of the right edge
        original_image_widths - nms_boxes_bs[:, :, 3],
        nms_boxes_bs[:, :, 0],
        # the mirrored location of the right edge is the image width
        # minus the position of the left edge
        original_image_widths - nms_boxes_bs[:, :, 1],
        nms_boxes_bs[:, :, 2],
        nms_scores_bs,
        nms_classes_bs,
    ]
  else:
    detections_bs = [
        image_ids_bs * t.ones_like(nms_scores_bs),
        nms_boxes_bs[:, :, 1],
        nms_boxes_bs[:, :, 0],
        nms_boxes_bs[:, :, 3],
        nms_boxes_bs[:, :, 2],
        nms_scores_bs,
        nms_classes_bs,
    ]
  return t.stack(detections_bs, axis=-1, name='detections')
//  [python]
def generate_detections_from_nms_output(nms_boxes_bs,
                                        nms_classes_bs,
                                        nms_scores_bs,
                                        image_ids,
                                        original_image_widths=None,
                                        flip=False):
  """Generating [id, x, y, w, h, score, class] from NMS outputs."""

  image_ids_bs = tf.cast(tf.expand_dims(image_ids, -1), nms_scores_bs.dtype)
  if flip:
    detections_bs = [
        image_ids_bs * tf.ones_like(nms_scores_bs),
        # the mirrored location of the left edge is the image width
        # minus the position of the right edge
        original_image_widths - nms_boxes_bs[:, :, 3],
        nms_boxes_bs[:, :, 0],
        # the mirrored location of the right edge is the image width
        # minus the position of the left edge
        original_image_widths - nms_boxes_bs[:, :, 1],
        nms_boxes_bs[:, :, 2],
        nms_scores_bs,
        nms_classes_bs,
    ]
  else:
    detections_bs = [
        image_ids_bs * tf.ones_like(nms_scores_bs),
        nms_boxes_bs[:, :, 1],
        nms_boxes_bs[:, :, 0],
        nms_boxes_bs[:, :, 3],
        nms_boxes_bs[:, :, 2],
        nms_scores_bs,
        nms_classes_bs,
    ]
  return tf.stack(detections_bs, axis=-1, name='detections')
//  [python]
def denormalize_points_by_image_size(points, image_sizes):
  """Denormalizes point coordinates by image sizes.

  Args:
    points: A tensor for normalized points by image size. Shape = [...,
      num_points, point_dim].
    image_sizes: A tensor for image sizes. Shape = [..., point_dim].

  Returns:
    A tensor for denormalized points. Shape = [..., num_points, point_dim].
  """
  if len(image_sizes.shape.as_list()) != len(points.shape.as_list()) - 1:
    raise ValueError(
        'Rank of `image_size` must be that of `points` minus 1: %d vs. %d.' %
        (len(image_sizes.shape.as_list()), len(points.shape.as_list())))
  return points * tf.expand_dims(
      tf.to_float(image_sizes), axis=-2)
//  [python]
def denormalize_points_by_image_size(points, image_sizes):
  """Denormalizes point coordinates by image sizes.

  Args:
    points: A tensor for normalized points by image size. Shape = [...,
      num_points, point_dim].
    image_sizes: A tensor for image sizes. Shape = [..., point_dim].

  Returns:
    A tensor for denormalized points. Shape = [..., num_points, point_dim].
  """
  if len(image_sizes.shape.as_list()) != len(points.shape.as_list()) - 1:
    raise ValueError(
        'Rank of `image_size` must be that of `points` minus 1: %d vs. %d.' %
        (len(image_sizes.shape.as_list()), len(points.shape.as_list())))
  return points * tf.expand_dims(
      tf.cast(image_sizes, dtype=points.dtype), axis=-2)
//  [python]
def preprocess_frame(frame):

    image = tf.convert_to_tensor(frame, dtype=tf.float32)
    image_gray = tf.image.rgb_to_grayscale(image)
    image_crop = tf.image.crop_to_bounding_box(image_gray, 34, 0, 160, 160)
    image_resize = tf.image.resize(image_crop, [84, 84])
    image_scaled = tf.divide(image_resize, 255)

    frame = image_scaled.numpy()[:, :, 0]

    return frame
//  [python]
def preprocess_frame(frame):

    image = tf.cast(tf.convert_to_tensor(frame), tf.float32)
    image_gray = tf.image.rgb_to_grayscale(image)
    image_crop = tf.image.crop_to_bounding_box(image_gray, 34, 0, 160, 160)
    image_resize = tf.image.resize(image_crop, [84, 84])
    image_scaled = tf.divide(image_resize, 255)

    frame = image_scaled.numpy()[:, :, 0]

    return frame
//  [python]
def naive_normalize_points(points, point_masks):
  """Naively normalizes points by shifting and scaling.

  Args:
    points: A tensor for points. Shape = [..., num_points, point_dim].
    point_masks: A tensor for point validities. Shape = [..., num_points].

  Returns:
    points: A tensor for normalized points. Shape = [..., num_points,
      point_dim].
  """
  point_masks = t.cast(point_masks, dtype=t.bool)
  point_dim = tf.shape(points)[-1]

  def compute_centers(points, point_masks):
    """Computes centers of valid points."""
    valid_points = tf.boolean_mask(points, point_masks)
    return tf.math.reduce_mean(valid_points, axis=-2, keepdims=True)

  def compute_max_spans(points, point_masks):
    """Computes maximum point set spans in any direction."""
    num_points = tf.shape(points)[-2]
    point_masks = data_utils.tile_last_dims(
        tf.expand_dims(tf.expand_dims(point_masks, axis=-1), axis=-1),
        last_dim_multiples=[num_points, point_dim])
    diffs = tf.math.abs(
        data_utils.tile_last_dims(
            tf.expand_dims(points, axis=-2), last_dim_multiples=[num_points, 1])
        - data_utils.tile_last_dims(
            tf.expand_dims(points, axis=-3),
            last_dim_multiples=[num_points, 1, 1]))
    diffs = tf.where(point_masks, diffs, tf.zeros_like(diffs))
    max_spans = tf.squeeze(
        tf.math.reduce_max(diffs, axis=[-3, -2, -1], keepdims=True), axis=[-2])
    return max_spans

  centers = compute_centers(points, point_masks)
  max_spans = compute_max_spans(points, point_masks)
  points = (points - centers) / tf.math.maximum(1e-12, max_spans)
  points = tf.where(
      data_utils.tile_last_dims(
          tf.expand_dims(point_masks, axis=-1), last_dim_multiples=[point_dim]),
      points, tf.zeros_like(points))
  return points
//  [python]
def naive_normalize_points(points, point_masks):
  """Naively normalizes points by shifting and scaling.

  Args:
    points: A tensor for points. Shape = [..., num_points, point_dim].
    point_masks: A tensor for point validities. Shape = [..., num_points].

  Returns:
    points: A tensor for normalized points. Shape = [..., num_points,
      point_dim].
  """
  point_masks = tf.cast(point_masks, dtype=tf.bool)
  point_dim = tf.shape(points)[-1]

  def compute_centers(points, point_masks):
    """Computes centers of valid points."""
    valid_points = tf.boolean_mask(points, point_masks)
    return tf.math.reduce_mean(valid_points, axis=-2, keepdims=True)

  def compute_max_spans(points, point_masks):
    """Computes maximum point set spans in any direction."""
    num_points = tf.shape(points)[-2]
    point_masks = data_utils.tile_last_dims(
        tf.expand_dims(tf.expand_dims(point_masks, axis=-1), axis=-1),
        last_dim_multiples=[num_points, point_dim])
    diffs = tf.math.abs(
        data_utils.tile_last_dims(
            tf.expand_dims(points, axis=-2), last_dim_multiples=[num_points, 1])
        - data_utils.tile_last_dims(
            tf.expand_dims(points, axis=-3),
            last_dim_multiples=[num_points, 1, 1]))
    diffs = tf.where(point_masks, diffs, tf.zeros_like(diffs))
    max_spans = tf.squeeze(
        tf.math.reduce_max(diffs, axis=[-3, -2, -1], keepdims=True), axis=[-2])
    return max_spans

  centers = compute_centers(points, point_masks)
  max_spans = compute_max_spans(points, point_masks)
  points = (points - centers) / tf.math.maximum(1e-12, max_spans)
  points = tf.where(
      data_utils.tile_last_dims(
          tf.expand_dims(point_masks, axis=-1), last_dim_multiples=[point_dim]),
      points, tf.zeros_like(points))
  return points
//  [python]
def _decode_areas(self, parsed_tensors):
    xmin = parsed_tensors['image/object/bbox/xmin']
    xmax = parsed_tensors['image/object/bbox/xmax']
    ymin = parsed_tensors['image/object/bbox/ymin']
    ymax = parsed_tensors['image/object/bbox/ymax']
    height = t.cast(parsed_tensors['image/height'], dtype=t.float32)
    width = t.cast(parsed_tensors['image/width'], dtype=t.float32)
    return tf.cond(
        tf.greater(tf.shape(parsed_tensors['image/object/area'])[0], 0),
        lambda: parsed_tensors['image/object/area'],
        lambda: (xmax - xmin) * (ymax - ymin) * height * width)
//  [python]
def _decode_areas(self, parsed_tensors):
    xmin = parsed_tensors['image/object/bbox/xmin']
    xmax = parsed_tensors['image/object/bbox/xmax']
    ymin = parsed_tensors['image/object/bbox/ymin']
    ymax = parsed_tensors['image/object/bbox/ymax']
    height = tf.cast(parsed_tensors['image/height'], dtype=tf.float32)
    width = tf.cast(parsed_tensors['image/width'], dtype=tf.float32)
    return tf.cond(
        tf.greater(tf.shape(parsed_tensors['image/object/area'])[0], 0),
        lambda: parsed_tensors['image/object/area'],
        lambda: (xmax - xmin) * (ymax - ymin) * height * width)
//  [python]
def _decode_masks(self, parsed_tensors):
    """Decode a set of PNG masks to the tf.float32 tensors."""

    def _decode_png_mask(png_bytes):
      mask = tf.squeeze(
          tf.io.decode_png(png_bytes, channels=1, dtype=tf.uint8), axis=-1)
      mask = tf.to_float(mask)
      mask.set_shape([None, None])
      return mask

    height = parsed_tensors['image/height']
    width = parsed_tensors['image/width']
    masks = parsed_tensors['image/object/mask']
    return tf.cond(
        pred=tf.greater(tf.size(input=masks), 0),
        true_fn=lambda: tf.map_fn(_decode_png_mask, masks, dtype=tf.float32),
        false_fn=lambda: tf.zeros([0, height, width], dtype=tf.float32))
//  [python]
def _decode_masks(self, parsed_tensors):
    """Decode a set of PNG masks to the tf.float32 tensors."""

    def _decode_png_mask(png_bytes):
      mask = tf.squeeze(
          tf.io.decode_png(png_bytes, channels=1, dtype=tf.uint8), axis=-1)
      mask = tf.cast(mask, dtype=tf.float32)
      mask.set_shape([None, None])
      return mask

    height = parsed_tensors['image/height']
    width = parsed_tensors['image/width']
    masks = parsed_tensors['image/object/mask']
    return tf.cond(
        pred=tf.greater(tf.size(input=masks), 0),
        true_fn=lambda: tf.map_fn(_decode_png_mask, masks, dtype=tf.float32),
        false_fn=lambda: tf.zeros([0, height, width], dtype=tf.float32))
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100 affect the loss
            active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 affect the loss
        loss_mask = tf.to_float(tf.cast(labels != -100, dtype=tf.bool))
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100 affect the loss
            active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 affect the loss
        loss_mask = tf.cast(labels != -100, dtype=unmasked_loss.dtype)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def build_reward(self): # reorder input % tour and return tour length (euclidean distance)
        self.permutations = tf.stack([tf.tile(tf.expand_dims(tf.range(self.batch_size,dtype=tf.int32),1),[1,self.max_length+1]),self.tour],2)
        if self.is_training==True:
            self.ordered_input_ = tf.gather_nd(self.input_,self.permutations)
        else:
            self.ordered_input_ = tf.gather_nd(tf.tile(self.input_,[self.batch_size,1,1]),self.permutations)
        self.ordered_input_ = tf.transpose(self.ordered_input_,[2,1,0]) # [features, seq length +1, batch_size]   Rq: +1 because end = start    
        
        ordered_x_ = self.ordered_input_[0] # ordered x, y coordinates [seq length +1, batch_size]
        ordered_y_ = self.ordered_input_[1] # ordered y coordinates [seq length +1, batch_size]          
        delta_x2 = tf.transpose(tf.square(ordered_x_[1:]-ordered_x_[:-1]),[1,0]) # [batch_size, seq length]        delta_x**2
        delta_y2 = tf.transpose(tf.square(ordered_y_[1:]-ordered_y_[:-1]),[1,0]) # [batch_size, seq length]        delta_y**2

        inter_city_distances = tf.sqrt(delta_x2+delta_y2) # sqrt(delta_x**2 + delta_y**2) this is the euclidean distance between each city: depot --> ... ---> depot      [batch_size, seq length]
        self.distances = tf.reduce_sum(inter_city_distances, axis=1) # [batch_size]
        self.reward = tf.to_float(self.distances) # define reward from tour length  
        tf.summary.scalar('reward_mean', tf.reduce_mean(self.reward))
//  [python]
def build_reward(self): # reorder input % tour and return tour length (euclidean distance)
        self.permutations = tf.stack([tf.tile(tf.expand_dims(tf.range(self.batch_size,dtype=tf.int32),1),[1,self.max_length+1]),self.tour],2)
        if self.is_training==True:
            self.ordered_input_ = tf.gather_nd(self.input_,self.permutations)
        else:
            self.ordered_input_ = tf.gather_nd(tf.tile(self.input_,[self.batch_size,1,1]),self.permutations)
        self.ordered_input_ = tf.transpose(self.ordered_input_,[2,1,0]) # [features, seq length +1, batch_size]   Rq: +1 because end = start    
        
        ordered_x_ = self.ordered_input_[0] # ordered x, y coordinates [seq length +1, batch_size]
        ordered_y_ = self.ordered_input_[1] # ordered y coordinates [seq length +1, batch_size]          
        delta_x2 = tf.transpose(tf.square(ordered_x_[1:]-ordered_x_[:-1]),[1,0]) # [batch_size, seq length]        delta_x**2
        delta_y2 = tf.transpose(tf.square(ordered_y_[1:]-ordered_y_[:-1]),[1,0]) # [batch_size, seq length]        delta_y**2

        inter_city_distances = tf.sqrt(delta_x2+delta_y2) # sqrt(delta_x**2 + delta_y**2) this is the euclidean distance between each city: depot --> ... ---> depot      [batch_size, seq length]
        self.distances = tf.reduce_sum(inter_city_distances, axis=1) # [batch_size]
        self.reward = tf.cast(self.distances,tf.float32) # define reward from tour length  
        tf.summary.scalar('reward_mean', tf.reduce_mean(self.reward))