jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def evaluate(inp_sentence):
  

  en_START_TOKEN_id = en_dictionary['<sos>']
  en_END_TOKEN_id   = en_dictionary['<eos>']
    
  pt_START_TOKEN_id = pt_dictionary['<sos>']
  pt_END_TOKEN_id   = pt_dictionary['<eos>']

  start_token = [pt_START_TOKEN_id]
  end_token   = [pt_END_TOKEN_id]

  
  # inp sentence is portuguese, hence adding the start and end token
  inp_sentence = start_token + encode(inp_sentence, pt_dictionary) + end_token
  encoder_input = tf.expand_dims(inp_sentence, 0)
  
  # as the target is english, the first word to the transformer should be the
  # english start token.
  decoder_input = [en_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 add_evaluation_step(result_tensor, ground_truth_tensor):
  """Inserts the operations we need to evaluate the accuracy of our results.

  Args:
    result_tensor: The new final node that produces results.
    ground_truth_tensor: The node we feed ground truth data
    into.

  Returns:
    Tuple of (evaluation step, prediction).
  """
  with tf.name_scope('accuracy'):
    with tf.name_scope('correct_prediction'):
      prediction = tf.argmax(result_tensor, 1)
      correct_prediction = tf.equal(prediction, ground_truth_tensor)
    with tf.name_scope('accuracy'):
      evaluation_step = tf.reduce_mean(tf.to_float(correct_prediction))
  tf.summary.scalar('accuracy', evaluation_step)
  return evaluation_step, prediction
//  [python]
def add_evaluation_step(result_tensor, ground_truth_tensor):
  """Inserts the operations we need to evaluate the accuracy of our results.

  Args:
    result_tensor: The new final node that produces results.
    ground_truth_tensor: The node we feed ground truth data
    into.

  Returns:
    Tuple of (evaluation step, prediction).
  """
  with tf.name_scope('accuracy'):
    with tf.name_scope('correct_prediction'):
      prediction = tf.argmax(result_tensor, 1)
      correct_prediction = tf.equal(prediction, ground_truth_tensor)
    with tf.name_scope('accuracy'):
      evaluation_step = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
  tf.summary.scalar('accuracy', evaluation_step)
  return evaluation_step, prediction
//  [python]
def _eval_per_sample(self, prob_flatten, cross_entropy_flatten, label_flatten):
        # Calculate the loss of each sample for hard negative mining
        loss_per_pixel = cross_entropy_flatten
        loss_per_pixel_zeros = tf.zeros_like(loss_per_pixel)
        loss_per_pixel_valid = tf.where(
            tf.less_equal(label_flatten, self.cfg.num_classes - 1),
            loss_per_pixel, 
            loss_per_pixel_zeros
        )
        loss_per_pixel_reshape = tf.reshape(loss_per_pixel_valid, [self.cfg.batch_size, self.cfg.input_size * self.cfg.input_size])
        loss_per_sample = tf.reduce_mean(t.cast(loss_per_pixel_reshape, tf.float32), axis=1)
        probs_per_pixel_reshape = tf.reshape(prob_flatten[:, :], [self.cfg.batch_size, self.cfg.input_size * self.cfg.input_size, self.cfg.num_classes])
        probs_per_sample = tf.reduce_mean(t.cast(probs_per_pixel_reshape, tf.float32), axis=1)

        return loss_per_sample, probs_per_sample
//  [python]
def _eval_per_sample(self, prob_flatten, cross_entropy_flatten, label_flatten):
        # Calculate the loss of each sample for hard negative mining
        loss_per_pixel = cross_entropy_flatten
        loss_per_pixel_zeros = tf.zeros_like(loss_per_pixel)
        loss_per_pixel_valid = tf.where(
            tf.less_equal(label_flatten, self.cfg.num_classes - 1),
            loss_per_pixel, 
            loss_per_pixel_zeros
        )
        loss_per_pixel_reshape = tf.reshape(loss_per_pixel_valid, [self.cfg.batch_size, self.cfg.input_size * self.cfg.input_size])
        loss_per_sample = tf.reduce_mean(tf.cast(loss_per_pixel_reshape, tf.float32), axis=1)
        probs_per_pixel_reshape = tf.reshape(prob_flatten[:, :], [self.cfg.batch_size, self.cfg.input_size * self.cfg.input_size, self.cfg.num_classes])
        probs_per_sample = tf.reduce_mean(tf.cast(probs_per_pixel_reshape, tf.float32), axis=1)

        return loss_per_sample, probs_per_sample
//  [python]
def preprocess_image(image_path, load_dims=False, read_mode="color"):
    global img_width, img_height, img_WIDTH, img_HEIGHT, aspect_ratio

    mode = "RGB" if read_mode == "color" else "L"
    channels = 3 if mode == 'RGB' else 1
    img = tf.io.read_file(image_path)  # Prevents crashes due to PNG images (ARGB)
    img = tf.image.decode_image(img, channels)
    img = tf.image.convert_image_dtype(img, tf.float32)

    img *= 255.

    shape = tf.cast(tf.shape(img)[:-1], tf.float32)

    if mode == "L":
        # Expand the 1 channel grayscale to 3 channel grayscale image
        temp = np.zeros(img.shape + (3,), dtype=np.uint8)
        temp[:, :, 0] = img
        temp[:, :, 1] = img.copy()
        temp[:, :, 2] = img.copy()

        img = temp

    if load_dims:
        img_WIDTH = img.shape[0]
        img_HEIGHT = img.shape[1]
        aspect_ratio = float(img_HEIGHT) / img_WIDTH

        img_width = args.img_size

        if maintain_aspect_ratio:
            img_height = int(img_width * aspect_ratio)
        else:
            img_height = args.img_size

    new_shape = tf.cast(tf.convert_to_tensor([img_width, img_height]), tf.int32)
    img = tf.image.resize(img, new_shape)

    img = img.numpy()

    # RGB -> BGR
    img = img[:, :, ::-1]

    img[:, :, 0] -= 103.939
    img[:, :, 1] -= 116.779
    img[:, :, 2] -= 123.68

    img = tf.convert_to_tensor(img, dtype=tf.float32)

    if K.image_data_format() == "channels_first":
        img = tf.cast(tf.transpose(img, (2, 0, 1)), tf.float32)

    img = tf.expand_dims(img, axis=0)
    return img
//  [python]
def _sum(self):
        with tf.name_scope('sum'):
            #  hist_input = tf.concat(
                #  [self.item_history_embedding, self.cate_history_embedding], 2
            #  )
            # add for cate
            hist_input = self.item_history_embedding
            self.target_item_embedding = self.target_item_embedding[:,:40]
            #  self.hist_embedding_sum = tf.reduce_sum(hist_input, 1)
            self.mask = self.iterator.mask
            self.real_mask = tf.to_float(self.mask)
            self.hist_embedding_sum = tf.reduce_sum(hist_input*tf.expand_dims(self.real_mask, -1), 1)

        model_output = tf.concat([self.target_item_embedding, self.hist_embedding_sum], -1)
        tf.summary.histogram("model_output", model_output)
        return model_output
//  [python]
def _sum(self):
        with tf.name_scope('sum'):
            #  hist_input = tf.concat(
                #  [self.item_history_embedding, self.cate_history_embedding], 2
            #  )
            # add for cate
            hist_input = self.item_history_embedding
            self.target_item_embedding = self.target_item_embedding[:,:40]
            #  self.hist_embedding_sum = tf.reduce_sum(hist_input, 1)
            self.mask = self.iterator.mask
            self.real_mask = tf.cast(self.mask, tf.float32)
            self.hist_embedding_sum = tf.reduce_sum(hist_input*tf.expand_dims(self.real_mask, -1), 1)

        model_output = tf.concat([self.target_item_embedding, self.hist_embedding_sum], -1)
        tf.summary.histogram("model_output", model_output)
        return model_output
//  [python]
def add_jpeg_decoding(input_width, input_height, input_depth, input_mean,
                      input_std):
  """Adds operations that perform JPEG decoding and resizing to the graph..

  Args:
    input_width: Desired width of the image fed into the recognizer graph.
    input_height: Desired width of the image fed into the recognizer graph.
    input_depth: Desired channels of the image fed into the recognizer graph.
    input_mean: Pixel value that should be zero in the image for the graph.
    input_std: How much to divide the pixel values by before recognition.

  Returns:
    Tensors for the node to feed JPEG data into, and the output of the
      preprocessing steps.
  """
  jpeg_data = tf.placeholder(tf.string, name='DecodeJPGInput')
  decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth)
  decoded_image_as_float = tf.to_float(decoded_image)
  decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
  resize_shape = tf.stack([input_height, input_width])
  resize_shape_as_int = tf.cast(resize_shape, dtype=tf.int32)
  resized_image = tf.image.resize_bilinear(decoded_image_4d,
                                           resize_shape_as_int)
  offset_image = tf.subtract(resized_image, input_mean)
  mul_image = tf.multiply(offset_image, 1.0 / input_std)
  return jpeg_data, mul_image
//  [python]
def add_jpeg_decoding(input_width, input_height, input_depth, input_mean,
                      input_std):
  """Adds operations that perform JPEG decoding and resizing to the graph..

  Args:
    input_width: Desired width of the image fed into the recognizer graph.
    input_height: Desired width of the image fed into the recognizer graph.
    input_depth: Desired channels of the image fed into the recognizer graph.
    input_mean: Pixel value that should be zero in the image for the graph.
    input_std: How much to divide the pixel values by before recognition.

  Returns:
    Tensors for the node to feed JPEG data into, and the output of the
      preprocessing steps.
  """
  jpeg_data = tf.placeholder(tf.string, name='DecodeJPGInput')
  decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth)
  decoded_image_as_float = tf.cast(decoded_image, dtype=tf.float32)
  decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
  resize_shape = tf.stack([input_height, input_width])
  resize_shape_as_int = tf.cast(resize_shape, dtype=tf.int32)
  resized_image = tf.image.resize_bilinear(decoded_image_4d,
                                           resize_shape_as_int)
  offset_image = tf.subtract(resized_image, input_mean)
  mul_image = tf.multiply(offset_image, 1.0 / input_std)
  return jpeg_data, mul_image
//  [python]
def causal_attention_mask(self, nd, ns, dtype):
        # 下半角矩阵
        i = tf.range(nd)[:, None]
        j = tf.range(ns)
        m = i > j - ns + nd
        return tf.to_float(m)
//  [python]
def causal_attention_mask(self, nd, ns, dtype):
        # 下半角矩阵
        i = tf.range(nd)[:, None]
        j = tf.range(ns)
        m = i > j - ns + nd
        return tf.cast(m, dtype)
//  [python]
def box_refinement_graph(box, gt_box, dtype_str):
    """Compute refinement needed to transform box to gt_box.
    box and gt_box are [N, (y1, x1, y2, x2)]
    """
    tensor_dtype = tf.float32

    box = tf.cast(box, tensor_dtype)
    gt_box = tf.cast(gt_box, tensor_dtype)

    height = box[:, 2] - box[:, 0]
    width = box[:, 3] - box[:, 1]
    center_y = box[:, 0] + 0.5 * height
    center_x = box[:, 1] + 0.5 * width

    gt_height = gt_box[:, 2] - gt_box[:, 0]
    gt_width = gt_box[:, 3] - gt_box[:, 1]
    gt_center_y = gt_box[:, 0] + 0.5 * gt_height
    gt_center_x = gt_box[:, 1] + 0.5 * gt_width

    dy = (gt_center_y - center_y) / height
    dx = (gt_center_x - center_x) / width
    dh = tf.log(gt_height / height)
    dw = tf.log(gt_width / width)

    result = tf.stack([dy, dx, dh, dw], axis=1)
    return result
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=self.config.GRADIENT_CLIP_NORM)
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keep_dims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.to_float(tf.size(w))
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keep_dims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=self.config.GRADIENT_CLIP_NORM)
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keep_dims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.cast(tf.size(w), tf.float32)
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keep_dims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def _cross_replica_average(self, t, num_shards_per_group):
    '''Calculates the average value of input tensor across TPU replicas.'''
    num_shards = tpu_function.get_tpu_context().number_of_shards
    group_assignment = None
    if num_shards_per_group > 1:
      if num_shards % num_shards_per_group != 0:
        raise ValueError('num_shards: %d mod shards_per_group: %d, should be 0'
                         % (num_shards, num_shards_per_group))
      num_groups = num_shards // num_shards_per_group
      group_assignment = [[
          x for x in range(num_shards) if x // num_shards_per_group == y
      ] for y in range(num_groups)]
      return tf.tpu.cross_replica_sum(t, group_assignment) / t.dtype.max
    else:
      tf.logging.info('TpuBatchNormalization None')
      return tf.tpu.cross_replica_sum(t, group_assignment) / t.dtype.max
//  [python]
def _cross_replica_average(self, t, num_shards_per_group):
    '''Calculates the average value of input tensor across TPU replicas.'''
    num_shards = tpu_function.get_tpu_context().number_of_shards
    group_assignment = None
    if num_shards_per_group > 1:
      if num_shards % num_shards_per_group != 0:
        raise ValueError('num_shards: %d mod shards_per_group: %d, should be 0'
                         % (num_shards, num_shards_per_group))
      num_groups = num_shards // num_shards_per_group
      group_assignment = [[
          x for x in range(num_shards) if x // num_shards_per_group == y
      ] for y in range(num_groups)]
      return tf.tpu.cross_replica_sum(t, group_assignment) / tf.cast(
          num_shards_per_group, t.dtype)
    else:
      tf.logging.info('TpuBatchNormalization None')
      return tf.tpu.cross_replica_sum(t, group_assignment) / tf.cast(
          num_shards, t.dtype)
//  [python]
def build_learning_rate(
    initial_lr,
    global_step,
    steps_per_epoch=None,
    lr_decay_type='exponential',
    decay_factor=0.97,
    decay_epochs=2.4,
    total_steps=None,
    warmup_epochs=5,
    start_from_step=0,
):
  '''Build learning rate.'''
  lr_step = global_step + start_from_step
  if lr_decay_type == 'exponential':
    assert steps_per_epoch is not None
    decay_steps = steps_per_epoch * decay_epochs
    lr = tf.train.exponential_decay(
        initial_lr, lr_step, decay_steps, decay_factor, staircase=True)
  elif lr_decay_type == 'cosine':
    assert total_steps is not None
    lr = 0.5 * initial_lr * (
        1 + tf.cos(np.pi * tf.cast(lr_step, tf.float32) / total_steps))
  elif lr_decay_type == 'constant':
    lr = initial_lr
  else:
    assert False, 'Unknown lr_decay_type : %s' % lr_decay_type

  if warmup_epochs:
    tf.logging.info('Learning rate warmup_epochs: %d' % warmup_epochs)
    warmup_steps = int(warmup_epochs * steps_per_epoch)
    warmup_lr = (
        initial_lr * tf.cast(lr_step, tf.float32) / tf.cast(
            warmup_steps, tf.float32))
    lr = tf.cond(lr_step < warmup_steps, lambda: warmup_lr, lambda: lr)

  return lr
//  [python]
def label_dataset(worker_id, prediction_dir, shard_id, num_shards):
  def label_dst_parser(value):
    keys_to_features = {
        'probabilities': tf.FixedLenFeature([FLAGS.num_label_classes], tf.float32),
        'classes': tf.FixedLenFeature([], tf.int64),
    }
    parsed = tf.parse_single_example(value, keys_to_features)
    features = {}
    features['probabilities'] = tf.cast(
        tf.reshape(parsed['probabilities'], shape=[FLAGS.num_label_classes]), dtype=tf.float32)
    features['classes'] = tf.cast(
        tf.reshape(parsed['classes'], shape=[]), dtype=tf.int32)
    return features
  input_file = os.path.join(
      prediction_dir,
      'train-info-%.5d-of-%.5d-%.5d' % (shard_id, num_shards, worker_id))
  dst = tf.data.Dataset.list_files(input_file)
  def fetch_dataset(filename):
    buffer_size = 8 * 1024 * 1024
    dataset = tf.data.TFRecordDataset(filename, buffer_size=buffer_size)
    return dataset
  dst = dst.apply(
      tf.data.experimental.parallel_interleave(
          fetch_dataset, cycle_length=1))
  dst = dst.apply(
      tf.data.experimental.map_and_batch(
          label_dst_parser, batch_size=1,
          num_parallel_batches=16))
  dst = dst.prefetch(tf.data.experimental.AUTOTUNE)
  return dst
//  [python]
def create_attention_mask_from_input_mask(from_tensor, to_mask):
  """Create 3D attention mask from a 2D tensor mask.

  Args:
    from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...].
    to_mask: int32 Tensor of shape [batch_size, to_seq_length].

  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3])
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]

  to_shape = get_shape_list(to_mask, expected_rank=2)
  to_seq_length = to_shape[1]

  to_mask = tf.to_float(
      tf.reshape(to_mask, [batch_size, 1, to_seq_length]))

  # We don't assume that `from_tensor` is a mask (although it could be). We
  # don't actually care if we attend *from* padding tokens (only *to* padding)
  # tokens so we create a tensor of all ones.
  #
  # `broadcast_ones` = [batch_size, from_seq_length, 1]
  broadcast_ones = tf.ones(
      shape=[batch_size, from_seq_length, 1], dtype=tf.float32)

  # Here we broadcast along two dimensions to create the mask.
  mask = broadcast_ones * to_mask

  return mask
//  [python]
def create_attention_mask_from_input_mask(from_tensor, to_mask):
  """Create 3D attention mask from a 2D tensor mask.

  Args:
    from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...].
    to_mask: int32 Tensor of shape [batch_size, to_seq_length].

  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3])
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]

  to_shape = get_shape_list(to_mask, expected_rank=2)
  to_seq_length = to_shape[1]

  to_mask = tf.cast(
      tf.reshape(to_mask, [batch_size, 1, to_seq_length]), tf.float32)

  # We don't assume that `from_tensor` is a mask (although it could be). We
  # don't actually care if we attend *from* padding tokens (only *to* padding)
  # tokens so we create a tensor of all ones.
  #
  # `broadcast_ones` = [batch_size, from_seq_length, 1]
  broadcast_ones = tf.ones(
      shape=[batch_size, from_seq_length, 1], dtype=tf.float32)

  # Here we broadcast along two dimensions to create the mask.
  mask = broadcast_ones * to_mask

  return mask
//  [python]
def _get_momentum(self, iteration):
    """Get the momentum value."""
    momentum = self._get_hyper("momentum")
    momentum_start = self._get_hyper("momentum_start")
    momentum_warm_up_steps = t.cast(
        self._get_hyper("warmup_steps"), iteration.dtype)
    value = tf.cond(
        (iteration - momentum_warm_up_steps) <= 0,
        true_fn=lambda: (momentum_start +  # pylint: disable=g-long-lambda
                         (t.cast(iteration, momentum.dtype) *
                          (momentum - momentum_start) / t.cast(
                              momentum_warm_up_steps, momentum.dtype))),
        false_fn=lambda: momentum)
    return value
//  [python]
def _get_momentum(self, iteration):
    """Get the momentum value."""
    momentum = self._get_hyper("momentum")
    momentum_start = self._get_hyper("momentum_start")
    momentum_warm_up_steps = tf.cast(
        self._get_hyper("warmup_steps"), iteration.dtype)
    value = tf.cond(
        (iteration - momentum_warm_up_steps) <= 0,
        true_fn=lambda: (momentum_start +  # pylint: disable=g-long-lambda
                         (tf.cast(iteration, momentum.dtype) *
                          (momentum - momentum_start) / tf.cast(
                              momentum_warm_up_steps, momentum.dtype))),
        false_fn=lambda: momentum)
    return value
//  [python]
def update(self, unscaled_grads_vars):
        """
        dynamically update the loss scaling
        """

        if self.global_norm is None:
            self.global_norm = tf.linalg.global_norm(
                [g for g in unscaled_grads_vars if g is not None]
            )

        isfinite = tf.math.is_finite(self.global_norm)
        # int rep of isfinite
        isfinite_int = t.cast(isfinite, t.int64)

        # inc step cntr
        steps_since_rescale = self._steps_since_rescale.assign_add(1)

        # if overflow, inc overflow cntr
        overflows_since_rescale = self._overflows_since_rescale.assign_add(
            1 - isfinite_int
        )

        ratio = t.cast(overflows_since_rescale, t.float32) / t.cast(
            steps_since_rescale, t.float32
        )

        # decrease LS condition
        dec_cond = t.cast(
            tf.math.greater(ratio, self.overflow_tolerance), t.int64
        )

        # reset cntrs on dec_cond
        overflows_since_rescale = overflows_since_rescale.assign(
            overflows_since_rescale * (1 - dec_cond)
        )

        steps_since_rescale = steps_since_rescale.assign(
            steps_since_rescale * (1 - dec_cond)
        )

        # update the loss scaling factor on Inf/NaN
        ls_fac_ = t.cast(1 + dec_cond, self._current_loss_scale.dtype)
        current_loss_scale = self._current_loss_scale.assign(
            self._current_loss_scale / ls_fac_
        )

        # increase loss scaling if steps_since_rescale > threshold
        ls_fac_ = tf.math.maximum(
            steps_since_rescale - self.increment_period + 1, 0
        )

        # increase loss 2x
        current_loss_scale = current_loss_scale.assign(
            current_loss_scale * t.cast(ls_fac_ + 1, current_loss_scale.dtype)
        )

        current_loss_scale = current_loss_scale.assign(
            tf.clip_by_value(
                current_loss_scale, self.min_loss_scale, self.max_loss_scale
            )
        )

        # reset counters if > threshold
        steps_since_rescale = steps_since_rescale.assign(
            steps_since_rescale * (1 - ls_fac_)
        )
        overflows_since_rescale = overflows_since_rescale.assign(
            overflows_since_rescale * (1 - ls_fac_)
        )

        return (
            tf.group(
                current_loss_scale,
                overflows_since_rescale,
                steps_since_rescale,
            ),
            isfinite,
        )
//  [python]
def update(self, unscaled_grads_vars):
        """
        dynamically update the loss scaling
        """

        if self.global_norm is None:
            self.global_norm = tf.linalg.global_norm(
                [g for g in unscaled_grads_vars if g is not None]
            )

        isfinite = tf.math.is_finite(self.global_norm)
        # int rep of isfinite
        isfinite_int = tf.cast(isfinite, tf.int64)

        # inc step cntr
        steps_since_rescale = self._steps_since_rescale.assign_add(1)

        # if overflow, inc overflow cntr
        overflows_since_rescale = self._overflows_since_rescale.assign_add(
            1 - isfinite_int
        )

        ratio = tf.cast(overflows_since_rescale, tf.float32) / tf.cast(
            steps_since_rescale, tf.float32
        )

        # decrease LS condition
        dec_cond = tf.cast(
            tf.math.greater(ratio, self.overflow_tolerance), tf.int64
        )

        # reset cntrs on dec_cond
        overflows_since_rescale = overflows_since_rescale.assign(
            overflows_since_rescale * (1 - dec_cond)
        )

        steps_since_rescale = steps_since_rescale.assign(
            steps_since_rescale * (1 - dec_cond)
        )

        # update the loss scaling factor on Inf/NaN
        ls_fac_ = tf.cast(1 + dec_cond, self._current_loss_scale.dtype)
        current_loss_scale = self._current_loss_scale.assign(
            self._current_loss_scale / ls_fac_
        )

        # increase loss scaling if steps_since_rescale > threshold
        ls_fac_ = tf.math.maximum(
            steps_since_rescale - self.increment_period + 1, 0
        )

        # increase loss 2x
        current_loss_scale = current_loss_scale.assign(
            current_loss_scale * tf.cast(ls_fac_ + 1, current_loss_scale.dtype)
        )

        current_loss_scale = current_loss_scale.assign(
            tf.clip_by_value(
                current_loss_scale, self.min_loss_scale, self.max_loss_scale
            )
        )

        # reset counters if > threshold
        steps_since_rescale = steps_since_rescale.assign(
            steps_since_rescale * (1 - ls_fac_)
        )
        overflows_since_rescale = overflows_since_rescale.assign(
            overflows_since_rescale * (1 - ls_fac_)
        )

        return (
            tf.group(
                current_loss_scale,
                overflows_since_rescale,
                steps_since_rescale,
            ),
            isfinite,
        )
//  [python]
def __init__(self,
               weight_decay=0.0,
               learning_rate=0.01,
               momentum=0.0,
               momentum_start=0.0,
               warmup_steps=1000,
               nesterov=False,
               name="SGD",
               weight_keys=("kernel", "weight"),
               bias_keys=("bias", "beta"),
               **kwargs):
    super(SGDTorch, self).__init__(name, **kwargs)

    # Create Hyper Params for each group of the LR
    self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
    self._set_hyper("bias_learning_rate", kwargs.get("lr", learning_rate))
    self._set_hyper("other_learning_rate", kwargs.get("lr", learning_rate))

    # SGD decay param
    self._set_hyper("decay", self._initial_decay)

    # Weight decay param
    self._weight_decay = weight_decay != 0.0
    self._set_hyper("weight_decay", weight_decay)

    # Enable Momentum
    self._momentum = False
    if isinstance(momentum, t.Tensor) or callable(momentum) or momentum > 0:
      self._momentum = True
    if isinstance(momentum, (int, float)) and (momentum < 0 or momentum > 1):
      raise ValueError("`momentum` must be between [0, 1].")
    self._set_hyper("momentum", momentum)
    self._set_hyper("momentum_start", momentum_start)
    self._set_hyper("warmup_steps", t.cast(warmup_steps, t.int32))

    # Enable Nesterov Momentum
    self.nesterov = nesterov

    # weights, biases, other
    self._weight_keys = weight_keys
    self._bias_keys = bias_keys
    self._variables_set = False
    self._wset = set()
    self._bset = set()
    self._oset = set()

    logging.info("Pytorch SGD simulation: ")
    logging.info("Weight Decay: %f", weight_decay)
//  [python]
def __init__(self,
               weight_decay=0.0,
               learning_rate=0.01,
               momentum=0.0,
               momentum_start=0.0,
               warmup_steps=1000,
               nesterov=False,
               name="SGD",
               weight_keys=("kernel", "weight"),
               bias_keys=("bias", "beta"),
               **kwargs):
    super(SGDTorch, self).__init__(name, **kwargs)

    # Create Hyper Params for each group of the LR
    self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
    self._set_hyper("bias_learning_rate", kwargs.get("lr", learning_rate))
    self._set_hyper("other_learning_rate", kwargs.get("lr", learning_rate))

    # SGD decay param
    self._set_hyper("decay", self._initial_decay)

    # Weight decay param
    self._weight_decay = weight_decay != 0.0
    self._set_hyper("weight_decay", weight_decay)

    # Enable Momentum
    self._momentum = False
    if isinstance(momentum, tf.Tensor) or callable(momentum) or momentum > 0:
      self._momentum = True
    if isinstance(momentum, (int, float)) and (momentum < 0 or momentum > 1):
      raise ValueError("`momentum` must be between [0, 1].")
    self._set_hyper("momentum", momentum)
    self._set_hyper("momentum_start", momentum_start)
    self._set_hyper("warmup_steps", tf.cast(warmup_steps, tf.int32))

    # Enable Nesterov Momentum
    self.nesterov = nesterov

    # weights, biases, other
    self._weight_keys = weight_keys
    self._bias_keys = bias_keys
    self._variables_set = False
    self._wset = set()
    self._bset = set()
    self._oset = set()

    logging.info("Pytorch SGD simulation: ")
    logging.info("Weight Decay: %f", weight_decay)
//  [python]
def _prepare_local(self, var_device, var_dtype, apply_state):
    super(SGDTorch, self)._prepare_local(var_device, var_dtype, apply_state)  # pytype: disable=attribute-error
    weight_decay = self._get_hyper("weight_decay")
    apply_state[(var_device,
                 var_dtype)]["weight_decay"] = t.cast(weight_decay, var_dtype)

    if self._momentum:
      momentum = self._get_momentum(self.iterations)
      momentum = t.cast(momentum, var_dtype)
      apply_state[(var_device,
                   var_dtype)]["momentum"] = t.identity(momentum)

    bias_lr = self._get_hyper("bias_learning_rate")
    if isinstance(bias_lr, LearningRateSchedule):
      bias_lr = bias_lr(self.iterations)
    bias_lr = t.cast(bias_lr, var_dtype)
    apply_state[(var_device,
                 var_dtype)]["bias_lr_t"] = t.identity(bias_lr)

    other_lr = self._get_hyper("other_learning_rate")
    if isinstance(other_lr, LearningRateSchedule):
      other_lr = other_lr(self.iterations)
    other_lr = t.cast(other_lr, var_dtype)
    apply_state[(var_device,
                 var_dtype)]["other_lr_t"] = t.identity(other_lr)

    return apply_state[(var_device, var_dtype)]
//  [python]
def _prepare_local(self, var_device, var_dtype, apply_state):
    super(SGDTorch, self)._prepare_local(var_device, var_dtype, apply_state)  # pytype: disable=attribute-error
    weight_decay = self._get_hyper("weight_decay")
    apply_state[(var_device,
                 var_dtype)]["weight_decay"] = tf.cast(weight_decay, var_dtype)

    if self._momentum:
      momentum = self._get_momentum(self.iterations)
      momentum = tf.cast(momentum, var_dtype)
      apply_state[(var_device,
                   var_dtype)]["momentum"] = tf.identity(momentum)

    bias_lr = self._get_hyper("bias_learning_rate")
    if isinstance(bias_lr, LearningRateSchedule):
      bias_lr = bias_lr(self.iterations)
    bias_lr = tf.cast(bias_lr, var_dtype)
    apply_state[(var_device,
                 var_dtype)]["bias_lr_t"] = tf.identity(bias_lr)

    other_lr = self._get_hyper("other_learning_rate")
    if isinstance(other_lr, LearningRateSchedule):
      other_lr = other_lr(self.iterations)
    other_lr = tf.cast(other_lr, var_dtype)
    apply_state[(var_device,
                 var_dtype)]["other_lr_t"] = tf.identity(other_lr)

    return apply_state[(var_device, var_dtype)]
//  [python]
def __init__(self, uri):
        """Create a new `BERTEncoder` object."""
        if not tf.test.is_gpu_available():
            glog.warning(
                "No GPU detected, BERT will run a lot slower than with a GPU.")

        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri, trainable=False)
            self._tokenizer = self._create_tokenizer_from_hub_module(uri)
            self._input_ids = tf.placeholder(
                name="input_ids", shape=[None, None], dtype=tf.int32)
            self._input_mask = tf.placeholder(
                name="input_mask", shape=[None, None], dtype=tf.int32)
            self._segment_ids = tf.zeros_like(self._input_ids)
            bert_inputs = dict(
                input_ids=self._input_ids,
                input_mask=self._input_mask,
                segment_ids=self._segment_ids
            )

            embeddings = embed_fn(
                inputs=bert_inputs, signature="tokens", as_dict=True)[
                "sequence_output"
            ]
            mask = tf.expand_dims(
                tf.to_float(self._input_mask), -1)
            self._embeddings = tf.reduce_sum(mask * embeddings, axis=1)

            init_ops = (
                tf.global_variables_initializer(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def __init__(self, uri):
        """Create a new `BERTEncoder` object."""
        if not tf.test.is_gpu_available():
            glog.warning(
                "No GPU detected, BERT will run a lot slower than with a GPU.")

        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri, trainable=False)
            self._tokenizer = self._create_tokenizer_from_hub_module(uri)
            self._input_ids = tf.placeholder(
                name="input_ids", shape=[None, None], dtype=tf.int32)
            self._input_mask = tf.placeholder(
                name="input_mask", shape=[None, None], dtype=tf.int32)
            self._segment_ids = tf.zeros_like(self._input_ids)
            bert_inputs = dict(
                input_ids=self._input_ids,
                input_mask=self._input_mask,
                segment_ids=self._segment_ids
            )

            embeddings = embed_fn(
                inputs=bert_inputs, signature="tokens", as_dict=True)[
                "sequence_output"
            ]
            mask = tf.expand_dims(
                tf.cast(self._input_mask, dtype=tf.float32), -1)
            self._embeddings = tf.reduce_sum(mask * embeddings, axis=1)

            init_ops = (
                tf.global_variables_initializer(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def scale(keypoints, y_scale, x_scale, scope=None):
  """Scales keypoint coordinates in x and y dimensions.

  Args:
    keypoints: a tensor of shape [num_instances, num_keypoints, 2]
    y_scale: (float) scalar tensor
    x_scale: (float) scalar tensor
    scope: name scope.

  Returns:
    new_keypoints: a tensor of shape [num_instances, num_keypoints, 2]
  """
  with tf.name_scope(scope, 'Scale'):
    y_scale = tf.to_float(y_scale)
    x_scale = tf.to_float(x_scale)
    new_keypoints = keypoints * [[[y_scale, x_scale]]]
    return new_keypoints
//  [python]
def scale(keypoints, y_scale, x_scale, scope=None):
  """Scales keypoint coordinates in x and y dimensions.

  Args:
    keypoints: a tensor of shape [num_instances, num_keypoints, 2]
    y_scale: (float) scalar tensor
    x_scale: (float) scalar tensor
    scope: name scope.

  Returns:
    new_keypoints: a tensor of shape [num_instances, num_keypoints, 2]
  """
  with tf.name_scope(scope, 'Scale'):
    y_scale = tf.cast(y_scale, tf.float32)
    x_scale = tf.cast(x_scale, tf.float32)
    new_keypoints = keypoints * [[[y_scale, x_scale]]]
    return new_keypoints
//  [python]
def create_attention_mask_from_keyword_mask(from_tensor, kw_mask_a, kw_mask_b, 
                                            real_mask_a, real_mask_b):
  """Create 3D keywords attention mask from two 2D tensor masks.
  A visualized version, see my evernote 20181026 2.4.2

  Args:
    from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...].
    kw_mask_a/b: int32 Tensor of shape [batch_size, to_seq_length], 
                 i.e. [1,0,1,0,0,0]. `1` means this position 
                      in input_ids is text_a/b's keyword
    real_mask_a/b: `1` means this position in input_ids is text_a/b's real token,
                   not CLS/SEP or padding token.
    ATTENTION: from_seq_length should be equal to to_seq_length, they all equal to 
               bert's sentence pair length (i.e. length of [cls]text_a[sep]text_b[sep])


  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3])
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]

  to_shape = get_shape_list(kw_mask_a, expected_rank=2)
  to_seq_length = to_shape[1]
  assert(from_seq_length==to_seq_length), \
        'Error: from_seq_length should be equal to to_seq_length'

  seq_length = from_seq_length

  kw_mask_a = tf.to_float(
      tf.reshape(kw_mask_a, [batch_size, 1, seq_length]))
  kw_mask_b = tf.to_float(
      tf.reshape(kw_mask_b, [batch_size, 1, seq_length]))
  real_mask_a = tf.to_float(
      tf.reshape(real_mask_a, [batch_size, seq_length, 1]))
  real_mask_b = tf.to_float(
      tf.reshape(real_mask_b, [batch_size, seq_length, 1]))
  kw_mask_cls = tf.math.add(kw_mask_a, kw_mask_b) # [B, 1, L]

  # We don't assume that `from_tensor` is a mask (although it could be). We
  # don't actually care if we attend *from* padding tokens (only *to* padding)
  # tokens so we create a tensor of all ones.
  #
  # `broadcast_ones` = [batch_size, seq_length, 1]
  broadcast_ones = tf.ones(
      shape=[batch_size, seq_length, 1], dtype=tf.float32)

  # Here we broadcast along two dimensions to create keyword mask.
  # kw_mask_a describes text_a's tokens attend to text_b's keywords, and vice versa.
  # kw_mask_a(i,j)=`1` means i'th token should attend to j'th token.
  # Thus, for kw_mask_a's elements, `1` only lies in (i,j) where i is text_a's token and j is text_b's keyword. other areas are all `0`.
  # Similar for kw_mask_b.
  # Then we add the two mask together to get final mask
  # Final mask size: [batch_size, seq_length, seq_length]
  kw_mask_a = broadcast_ones * real_mask_a * kw_mask_b # a attend b's keyword, 
  kw_mask_b = broadcast_ones * real_mask_b * kw_mask_a
  kw_mask = tf.math.add(kw_mask_a, kw_mask_b)

  return kw_mask
//  [python]
def create_attention_mask_from_keyword_mask(from_tensor, kw_mask_a, kw_mask_b, 
                                            real_mask_a, real_mask_b):
  """Create 3D keywords attention mask from two 2D tensor masks.
  A visualized version, see my evernote 20181026 2.4.2

  Args:
    from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...].
    kw_mask_a/b: int32 Tensor of shape [batch_size, to_seq_length], 
                 i.e. [1,0,1,0,0,0]. `1` means this position 
                      in input_ids is text_a/b's keyword
    real_mask_a/b: `1` means this position in input_ids is text_a/b's real token,
                   not CLS/SEP or padding token.
    ATTENTION: from_seq_length should be equal to to_seq_length, they all equal to 
               bert's sentence pair length (i.e. length of [cls]text_a[sep]text_b[sep])


  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3])
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]

  to_shape = get_shape_list(kw_mask_a, expected_rank=2)
  to_seq_length = to_shape[1]
  assert(from_seq_length==to_seq_length), \
        'Error: from_seq_length should be equal to to_seq_length'

  seq_length = from_seq_length

  kw_mask_a = tf.cast(
      tf.reshape(kw_mask_a, [batch_size, 1, seq_length]), tf.float32)
  kw_mask_b = tf.cast(
      tf.reshape(kw_mask_b, [batch_size, 1, seq_length]), tf.float32)
  real_mask_a = tf.cast(
      tf.reshape(real_mask_a, [batch_size, seq_length, 1]), tf.float32)
  real_mask_b = tf.cast(
      tf.reshape(real_mask_b, [batch_size, seq_length, 1]), tf.float32)
  kw_mask_cls = tf.math.add(kw_mask_a, kw_mask_b) # [B, 1, L]

  # We don't assume that `from_tensor` is a mask (although it could be). We
  # don't actually care if we attend *from* padding tokens (only *to* padding)
  # tokens so we create a tensor of all ones.
  #
  # `broadcast_ones` = [batch_size, seq_length, 1]
  broadcast_ones = tf.ones(
      shape=[batch_size, seq_length, 1], dtype=tf.float32)

  # Here we broadcast along two dimensions to create keyword mask.
  # kw_mask_a describes text_a's tokens attend to text_b's keywords, and vice versa.
  # kw_mask_a(i,j)=`1` means i'th token should attend to j'th token.
  # Thus, for kw_mask_a's elements, `1` only lies in (i,j) where i is text_a's token and j is text_b's keyword. other areas are all `0`.
  # Similar for kw_mask_b.
  # Then we add the two mask together to get final mask
  # Final mask size: [batch_size, seq_length, seq_length]
  kw_mask_a = broadcast_ones * real_mask_a * kw_mask_b # a attend b's keyword, 
  kw_mask_b = broadcast_ones * real_mask_b * kw_mask_a
  kw_mask = tf.math.add(kw_mask_a, kw_mask_b)

  return kw_mask
//  [python]
def _top_singular_vector(m):
        # handle the case where m is a tensor of rank 0 or rank 1.
        # Example:
        #   scalar (rank 0) a, shape []=> [[a]], shape [1,1]
        #   vector (rank 1) [a,b], shape [2] => [[a,b]], shape [1,2]
        original_rank = tf.rank(m)
        shape = tf.shape(m)
        first_pad = tf.to_int32(tf.less(original_rank, 2))
        second_pad = tf.to_int32(tf.equal(original_rank, 0))
        new_shape = tf.concat(
            [
                tf.ones(shape=first_pad, dtype=tf.int32),
                tf.ones(shape=second_pad, dtype=tf.int32),
                shape,
            ],
            axis=0,
        )
        n = tf.reshape(m, new_shape)
        st, ut, vt = tf.linalg.svd(n, full_matrices=False)
        n_size = tf.shape(n)
        ut = tf.reshape(ut[:, 0], [n_size[0], 1])
        vt = tf.reshape(vt[:, 0], [n_size[1], 1])
        st = tf.matmul(ut, tf.transpose(vt))
        # when we return the top singular vector, we have to remove the
        # dimension we have added on
        st_shape = tf.shape(st)
        begin = tf.to_int32(tf.less(original_rank, 2))
        end = 2 - tf.to_int32(tf.equal(original_rank, 0))
        new_shape = st_shape[begin:end]
        return tf.reshape(st, new_shape)
//  [python]
def _top_singular_vector(m):
        # handle the case where m is a tensor of rank 0 or rank 1.
        # Example:
        #   scalar (rank 0) a, shape []=> [[a]], shape [1,1]
        #   vector (rank 1) [a,b], shape [2] => [[a,b]], shape [1,2]
        original_rank = tf.rank(m)
        shape = tf.shape(m)
        first_pad = tf.cast(tf.less(original_rank, 2), dtype=tf.int32)
        second_pad = tf.cast(tf.equal(original_rank, 0), dtype=tf.int32)
        new_shape = tf.concat(
            [
                tf.ones(shape=first_pad, dtype=tf.int32),
                tf.ones(shape=second_pad, dtype=tf.int32),
                shape,
            ],
            axis=0,
        )
        n = tf.reshape(m, new_shape)
        st, ut, vt = tf.linalg.svd(n, full_matrices=False)
        n_size = tf.shape(n)
        ut = tf.reshape(ut[:, 0], [n_size[0], 1])
        vt = tf.reshape(vt[:, 0], [n_size[1], 1])
        st = tf.matmul(ut, tf.transpose(vt))
        # when we return the top singular vector, we have to remove the
        # dimension we have added on
        st_shape = tf.shape(st)
        begin = tf.cast(tf.less(original_rank, 2), dtype=tf.int32)
        end = 2 - tf.cast(tf.equal(original_rank, 0), dtype=tf.int32)
        new_shape = st_shape[begin:end]
        return tf.reshape(st, new_shape)