jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _build_train(self):
    print("-" * 80)
    print("Build train graph")
    logits = self._model(self.x_train, is_training=True)
    log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=self.y_train)
    self.loss = tf.reduce_mean(log_probs)

    if self.use_aux_heads:
      log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=self.aux_logits, labels=self.y_train)
      self.aux_loss = tf.reduce_mean(log_probs)
      train_loss = self.loss + 0.4 * self.aux_loss
    else:
      train_loss = self.loss

    self.train_preds = tf.argmax(logits, axis=1)
    self.train_preds = tf.cast(self.train_preds, tf.int32)
    self.train_acc = tf.equal(self.train_preds, self.y_train)
    self.train_acc = tf.cast(self.train_acc, tf.int32)
    self.train_acc = tf.reduce_sum(self.train_acc)

    tf_variables = [
      var for var in tf.trainable_variables() if (
        var.name.startswith(self.name) and "aux_head" not in var.name)]
    self.num_vars = count_model_params(tf_variables)
    print("Model has {0} params".format(self.num_vars))

    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      train_loss,
      tf_variables,
      self.global_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      lr_cosine=self.lr_cosine,
      lr_max=self.lr_max,
      lr_min=self.lr_min,
      lr_T_0=self.lr_T_0,
      lr_T_mul=self.lr_T_mul,
      num_train_batches=self.num_train_batches,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)
//  [python]
def laid_out_slice_num(self, tensor_shape):
    """A LaidOutTensor with an int32 scalar, identical for identical slices.

    This is useful for synchronizing random operations.

    Args:
      tensor_shape: a TensorShape
    Returns:
      a LaidOutTensor where each slice is an integer scalar.
    """
    ret = self.slicewise(lambda: tf.to_int32(0))
    tensor_layout = self.tensor_layout(tensor_shape)
    for mesh_axis in tensor_layout.tensor_axis_to_mesh_axis:
      if mesh_axis is not None:
        def my_fn(x, pcoord, mesh_dim_size):
          return x * mesh_dim_size + pcoord
        ret = self.slicewise(
            my_fn, ret, self.laid_out_pcoord(mesh_axis),
            self.shape[mesh_axis].size)
    return ret
//  [python]
def laid_out_slice_num(self, tensor_shape):
    """A LaidOutTensor with an int32 scalar, identical for identical slices.

    This is useful for synchronizing random operations.

    Args:
      tensor_shape: a TensorShape
    Returns:
      a LaidOutTensor where each slice is an integer scalar.
    """
    ret = self.slicewise(lambda: tf.cast(0, tf.int32))
    tensor_layout = self.tensor_layout(tensor_shape)
    for mesh_axis in tensor_layout.tensor_axis_to_mesh_axis:
      if mesh_axis is not None:
        def my_fn(x, pcoord, mesh_dim_size):
          return x * mesh_dim_size + pcoord
        ret = self.slicewise(
            my_fn, ret, self.laid_out_pcoord(mesh_axis),
            self.shape[mesh_axis].size)
    return ret
//  [python]
def _build_model(self, weights=None):
        """
        Builds TensorFlow model.
        :return:
        """
        # initialize weights and biases
        self._initialize_weights(weights)

        # TensorFlow operations
        self.visible_units_placeholder = tf.placeholder(tf.float32, shape=[None, self.n_visible_units])
        self.compute_hidden_units_op = self._activation_function_class(
            tf.transpose(tf.matmul(self.W, tf.transpose(self.visible_units_placeholder))) + self.c)
        self.hidden_units_placeholder = tf.placeholder(tf.float32, shape=[None, self.n_hidden_units])
        self.compute_visible_units_op = self._activation_function_class(
            tf.matmul(self.hidden_units_placeholder, self.W) + self.b)
        self.random_uniform_values = tf.Variable(tf.random_uniform([self.batch_size, self.n_hidden_units]))
        sample_hidden_units_op = tf.to_float(self.random_uniform_values < self.compute_hidden_units_op)
        self.random_variables = [self.random_uniform_values]

        # Positive gradient
        # Outer product. N is the batch size length.
        # From http://stackoverflow.com/questions/35213787/tensorflow-batch-outer-product
        positive_gradient_op = tf.matmul(tf.expand_dims(sample_hidden_units_op, 2),  # [N, U, 1]
                                         tf.expand_dims(self.visible_units_placeholder, 1))  # [N, 1, V]

        # Negative gradient
        # Gibbs sampling
        sample_hidden_units_gibbs_step_op = sample_hidden_units_op
        for t in range(self.contrastive_divergence_iter):
            compute_visible_units_op = self._activation_function_class(
                tf.matmul(sample_hidden_units_gibbs_step_op, self.W) + self.b)
            compute_hidden_units_gibbs_step_op = self._activation_function_class(
                tf.transpose(tf.matmul(self.W, tf.transpose(compute_visible_units_op))) + self.c)
            random_uniform_values = tf.Variable(tf.random_uniform([self.batch_size, self.n_hidden_units]))
            sample_hidden_units_gibbs_step_op = tf.to_float(random_uniform_values < compute_hidden_units_gibbs_step_op)
            self.random_variables.append(random_uniform_values)

        negative_gradient_op = tf.matmul(tf.expand_dims(sample_hidden_units_gibbs_step_op, 2),  # [N, U, 1]
                                         tf.expand_dims(compute_visible_units_op, 1))  # [N, 1, V]

        compute_delta_W = tf.reduce_mean(positive_gradient_op - negative_gradient_op, 0)
        compute_delta_b = tf.reduce_mean(self.visible_units_placeholder - compute_visible_units_op, 0)
        compute_delta_c = tf.reduce_mean(sample_hidden_units_op - sample_hidden_units_gibbs_step_op, 0)

        self.update_W = tf.assign_add(self.W, self.learning_rate * compute_delta_W)
        self.update_b = tf.assign_add(self.b, self.learning_rate * compute_delta_b)
        self.update_c = tf.assign_add(self.c, self.learning_rate * compute_delta_c)
//  [python]
def _build_model(self, weights=None):
        """
        Builds TensorFlow model.
        :return:
        """
        # initialize weights and biases
        self._initialize_weights(weights)

        # TensorFlow operations
        self.visible_units_placeholder = tf.placeholder(tf.float32, shape=[None, self.n_visible_units])
        self.compute_hidden_units_op = self._activation_function_class(
            tf.transpose(tf.matmul(self.W, tf.transpose(self.visible_units_placeholder))) + self.c)
        self.hidden_units_placeholder = tf.placeholder(tf.float32, shape=[None, self.n_hidden_units])
        self.compute_visible_units_op = self._activation_function_class(
            tf.matmul(self.hidden_units_placeholder, self.W) + self.b)
        self.random_uniform_values = tf.Variable(tf.random_uniform([self.batch_size, self.n_hidden_units]))
        sample_hidden_units_op = tf.cast(self.random_uniform_values < self.compute_hidden_units_op, tf.float32)
        self.random_variables = [self.random_uniform_values]

        # Positive gradient
        # Outer product. N is the batch size length.
        # From http://stackoverflow.com/questions/35213787/tensorflow-batch-outer-product
        positive_gradient_op = tf.matmul(tf.expand_dims(sample_hidden_units_op, 2),  # [N, U, 1]
                                         tf.expand_dims(self.visible_units_placeholder, 1))  # [N, 1, V]

        # Negative gradient
        # Gibbs sampling
        sample_hidden_units_gibbs_step_op = sample_hidden_units_op
        for t in range(self.contrastive_divergence_iter):
            compute_visible_units_op = self._activation_function_class(
                tf.matmul(sample_hidden_units_gibbs_step_op, self.W) + self.b)
            compute_hidden_units_gibbs_step_op = self._activation_function_class(
                tf.transpose(tf.matmul(self.W, tf.transpose(compute_visible_units_op))) + self.c)
            random_uniform_values = tf.Variable(tf.random_uniform([self.batch_size, self.n_hidden_units]))
            sample_hidden_units_gibbs_step_op = tf.cast(random_uniform_values < compute_hidden_units_gibbs_step_op, tf.float32)
            self.random_variables.append(random_uniform_values)

        negative_gradient_op = tf.matmul(tf.expand_dims(sample_hidden_units_gibbs_step_op, 2),  # [N, U, 1]
                                         tf.expand_dims(compute_visible_units_op, 1))  # [N, 1, V]

        compute_delta_W = tf.reduce_mean(positive_gradient_op - negative_gradient_op, 0)
        compute_delta_b = tf.reduce_mean(self.visible_units_placeholder - compute_visible_units_op, 0)
        compute_delta_c = tf.reduce_mean(sample_hidden_units_op - sample_hidden_units_gibbs_step_op, 0)

        self.update_W = tf.assign_add(self.W, self.learning_rate * compute_delta_W)
        self.update_b = tf.assign_add(self.b, self.learning_rate * compute_delta_b)
        self.update_c = tf.assign_add(self.c, self.learning_rate * compute_delta_c)
//  [python]
def create_input_queue(create_tensor_dict_fn):
  """Sets up reader, prefetcher and returns input queue.

  Args:
    create_tensor_dict_fn: function to create tensor dictionary.

  Returns:
    all_dict: A dictionary holds tensors for images, boxes, and targets.
  """
  tensor_dict = create_tensor_dict_fn()
  all_dict = {}

  num_images = len(tensor_dict[fields.InputDataFields.image])
  all_dict['batch'] = tensor_dict['batch']
  del tensor_dict['batch']

  for i in range(num_images):
    suffix = str(i)
    for key, val in tensor_dict.items():
      all_dict[key + suffix] = val[i]

    all_dict[fields.InputDataFields.image + suffix] = tf.to_float(
        tf.expand_dims(all_dict[fields.InputDataFields.image + suffix], 0))

  return all_dict
//  [python]
def create_input_queue(create_tensor_dict_fn):
  """Sets up reader, prefetcher and returns input queue.

  Args:
    create_tensor_dict_fn: function to create tensor dictionary.

  Returns:
    all_dict: A dictionary holds tensors for images, boxes, and targets.
  """
  tensor_dict = create_tensor_dict_fn()
  all_dict = {}

  num_images = len(tensor_dict[fields.InputDataFields.image])
  all_dict['batch'] = tensor_dict['batch']
  del tensor_dict['batch']

  for i in range(num_images):
    suffix = str(i)
    for key, val in tensor_dict.items():
      all_dict[key + suffix] = val[i]

    all_dict[fields.InputDataFields.image + suffix] = tf.cast(
        tf.expand_dims(all_dict[fields.InputDataFields.image + suffix], 0), tf.float32)

  return all_dict
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])
        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, is_training=True, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.to_int32(valid_shuffle_preds)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.to_int32(self.valid_shuffle_acc)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])
        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, is_training=True, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.cast(valid_shuffle_preds, tf.int32)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.cast(self.valid_shuffle_acc, tf.int32)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)
//  [python]
def embedding_to_padding(emb):
    """Calculates the padding mask based on which embeddings are all zero.

  We have hacked symbol_modality to return all-zero embeddings for padding.

  Args:
    emb: a Tensor with shape [..., depth].

  Returns:
    a float Tensor with shape [...]. Each element is 1 if its corresponding
      embedding vector is all zero, and is 0 otherwise.
  """
    emb_sum = tf.reduce_sum(tf.abs(emb), axis = -1)
    return tf.to_float(tf.equal(emb_sum, 0.0))
//  [python]
def embedding_to_padding(emb):
    """Calculates the padding mask based on which embeddings are all zero.

  We have hacked symbol_modality to return all-zero embeddings for padding.

  Args:
    emb: a Tensor with shape [..., depth].

  Returns:
    a float Tensor with shape [...]. Each element is 1 if its corresponding
      embedding vector is all zero, and is 0 otherwise.
  """
    emb_sum = tf.reduce_sum(tf.abs(emb), axis = -1)
    return tf.cast(tf.equal(emb_sum, 0.0), tf.float32)
//  [python]
def scale_dimension(dim, scale):
  """Scales the input dimension.

  Args:
    dim: Input dimension (a scalar).
    scale: The amount of scaling applied to the input.

  Returns:
    Scaled dimension.
  """
  if isinstance(dim, tf.Tensor):
    return tf.cast((tf.to_float(dim) - 1.0) * scale + 1.0, dtype=tf.int32)
  else:
    return int((float(dim) - 1.0) * scale + 1.0)
//  [python]
def scale_dimension(dim, scale):
  """Scales the input dimension.

  Args:
    dim: Input dimension (a scalar).
    scale: The amount of scaling applied to the input.

  Returns:
    Scaled dimension.
  """
  if isinstance(dim, tf.Tensor):
    return tf.cast((tf.cast(dim, tf.float32) - 1.0) * scale + 1.0, dtype=tf.int32)
  else:
    return int((float(dim) - 1.0) * scale + 1.0)
//  [python]
def bit_to_int(x_bit, num_bits, base = 2):
    """Turn x_bit representing numbers bitwise (lower-endian) to int tensor.

  Args:
    x_bit: Tensor containing numbers in a particular base to be converted to
      int.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Integer representation of this number.
  """
    x_l = tf.to_int64(tf.reshape(x_bit, [-1, num_bits]))
    x_labels = [
        x_l[:, i] * tf.to_int64(base) ** tf.to_int64(i) for i in range(num_bits)
    ]
    res = sum(x_labels)
    return tf.to_int64(tf.reshape(res, x_bit.get_shape().as_list()[:-1]))
//  [python]
def bit_to_int(x_bit, num_bits, base = 2):
    """Turn x_bit representing numbers bitwise (lower-endian) to int tensor.

  Args:
    x_bit: Tensor containing numbers in a particular base to be converted to
      int.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Integer representation of this number.
  """
    x_l = tf.cast(tf.reshape(x_bit, [-1, num_bits]), tf.int64)
    x_labels = [
        x_l[:, i] * tf.cast(base, tf.int64) ** tf.cast(i, tf.int64) for i in range(num_bits)
    ]
    res = sum(x_labels)
    return tf.cast(tf.reshape(res, x_bit.get_shape().as_list()[:-1]), tf.int64)
//  [python]
def get_ignore_mask(labels, flags):
  """Get the mask of elements to ignore.

  Args:
    labels: Ground truth label of shape [batch, height, width].
    flags: The input flags

  Returns:
    3D mask of same dimensionality as labels. 0 for elements to ignore, 1 for
    elements to keep.
  """
  mask = tf.to_int32(
      tf.not_equal(labels,
                   model_input.dataset_descriptors[flags.dataset].ignore_label))
  return mask
//  [python]
def get_ignore_mask(labels, flags):
  """Get the mask of elements to ignore.

  Args:
    labels: Ground truth label of shape [batch, height, width].
    flags: The input flags

  Returns:
    3D mask of same dimensionality as labels. 0 for elements to ignore, 1 for
    elements to keep.
  """
  mask = tf.cast(
      tf.not_equal(labels,
                   model_input.dataset_descriptors[flags.dataset].ignore_label), tf.int32)
  return mask
//  [python]
def get_learning_rate(learning_rate, hidden_size, learning_rate_warmup_steps):
  """Calculate learning rate with linear warmup and rsqrt decay."""
  with tf.name_scope("learning_rate"):
    warmup_steps = tf.to_float(learning_rate_warmup_steps)
    step = tf.to_float(tf.train.get_or_create_global_step())

    learning_rate *= (hidden_size ** -0.5)
    # Apply linear warmup
    learning_rate *= tf.minimum(1.0, step / warmup_steps)
    # Apply rsqrt decay
    learning_rate *= tf.rsqrt(tf.maximum(step, warmup_steps))

    # Create a named tensor that will be logged using the logging hook.
    # The full name includes variable and names scope. In this case, the name
    # is model/get_train_op/learning_rate/learning_rate
    tf.identity(learning_rate, "learning_rate")

    return learning_rate
//  [python]
def get_learning_rate(learning_rate, hidden_size, learning_rate_warmup_steps):
  """Calculate learning rate with linear warmup and rsqrt decay."""
  with tf.name_scope("learning_rate"):
    warmup_steps = tf.cast(learning_rate_warmup_steps, dtype=tf.float32)
    step = tf.cast(tf.train.get_or_create_global_step(), dtype=tf.float32)

    learning_rate *= (hidden_size ** -0.5)
    # Apply linear warmup
    learning_rate *= tf.minimum(1.0, step / warmup_steps)
    # Apply rsqrt decay
    learning_rate *= tf.rsqrt(tf.maximum(step, warmup_steps))

    # Create a named tensor that will be logged using the logging hook.
    # The full name includes variable and names scope. In this case, the name
    # is model/get_train_op/learning_rate/learning_rate
    tf.identity(learning_rate, "learning_rate")

    return learning_rate
//  [python]
def accuracy(self, y_pred, y_true, x_in):		# y_pred is [-1, self.out_seq_len, num_decoder_symbols]; y_true is [-1, self.out_seq_len]
        '''
        Compute accuracy of the prediction, based on the true labels.  Use the average number of equal
        values.
        '''
        pred_idx = tf.to_int32(tf.argmax(y_pred, 2))		# [-1, self.out_seq_len]
        if self.verbose > 2: print ("my_accuracy pred_idx = %s" % pred_idx)
        accuracy = tf.reduce_mean(tf.cast(tf.equal(pred_idx, y_true), tf.float32), name='acc')
        return accuracy
//  [python]
def accuracy(self, y_pred, y_true, x_in):		# y_pred is [-1, self.out_seq_len, num_decoder_symbols]; y_true is [-1, self.out_seq_len]
        '''
        Compute accuracy of the prediction, based on the true labels.  Use the average number of equal
        values.
        '''
        pred_idx = tf.cast(tf.argmax(y_pred, 2), tf.int32)		# [-1, self.out_seq_len]
        if self.verbose > 2: print ("my_accuracy pred_idx = %s" % pred_idx)
        accuracy = tf.reduce_mean(tf.cast(tf.equal(pred_idx, y_true), tf.float32), name='acc')
        return accuracy
//  [python]
def foreground_iou(labels, predictions, flags):
  """Calculates the iou of class 1 ignoring class 0."""
  weights = tf.to_float(
      tf.not_equal(labels,
                   model_input.dataset_descriptors[flags.dataset].ignore_label))

  predictions = tf.to_float(predictions)
  labels = tf.to_float(labels)
  labels *= weights
  predictions *= weights

  # Use this instead of built in mean_iou,
  # because we want to ignnore the background label.
  # This gives the IOU the same meaning as it does in our baseline
  # https://arxiv.org/pdf/1603.06180.pdf
  intersect = labels * predictions
  union = labels + predictions - intersect

  intersect_sum = tf.reduce_sum(intersect, [1, 2, 3])
  union_sum = tf.reduce_sum(union, [1, 2, 3])
  # Make sure we don't divide by 0.
  iou = tf.where(
      tf.less(union_sum, 1e-20), tf.ones_like(union_sum),
      intersect_sum / union_sum)

  return iou, weights
//  [python]
def foreground_iou(labels, predictions, flags):
  """Calculates the iou of class 1 ignoring class 0."""
  weights = tf.cast(
      tf.not_equal(labels,
                   model_input.dataset_descriptors[flags.dataset].ignore_label), tf.float32)

  predictions = tf.cast(predictions, tf.float32)
  labels = tf.cast(labels, tf.float32)
  labels *= weights
  predictions *= weights

  # Use this instead of built in mean_iou,
  # because we want to ignnore the background label.
  # This gives the IOU the same meaning as it does in our baseline
  # https://arxiv.org/pdf/1603.06180.pdf
  intersect = labels * predictions
  union = labels + predictions - intersect

  intersect_sum = tf.reduce_sum(intersect, [1, 2, 3])
  union_sum = tf.reduce_sum(union, [1, 2, 3])
  # Make sure we don't divide by 0.
  iou = tf.where(
      tf.less(union_sum, 1e-20), tf.ones_like(union_sum),
      intersect_sum / union_sum)

  return iou, weights
//  [python]
def apply_gradients(self, grads_and_vars, global_step, name=None):
    """Perform an update with the parameters."""

    # we meta-train with 10k steps. When applying to longer problems we want to
    # have a reasonable schedule so we rescale.

    rescale_global_step = float(10000) / self._training_steps * tf.to_float(
        global_step)

    beta1_power = tf.get_variable(
        dtype=tf.float32, name="beta1_power", initializer=self._beta1)
    beta2_power = tf.get_variable(
        dtype=tf.float32, name="beta2_power", initializer=self._beta2)

    exp_factor = tf.exp(-self._exponential_decay *
                        tf.to_float(rescale_global_step))

    # lr reduction per step.
    linear_factor = tf.maximum(
        1 - self._linear_decay * tf.to_float(rescale_global_step), 0.0)

    lr = exp_factor * linear_factor * self._learning_rate

    assignments = []
    for (grad, param) in grads_and_vars:
      if grad is None or param is None:
        continue
      # sparse to dense conversion
      grad = tf.convert_to_tensor(grad)

      param_name = self._get_variable_name(param.name)

      m = tf.get_variable(
          name=param_name + "/adam_m",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())
      v = tf.get_variable(
          name=param_name + "/adam_v",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())

      next_m = (self._beta1 * m + (1.0 - self._beta1) * grad)
      next_v = (self._beta2 * v + (1.0 - self._beta2) * tf.square(grad))
      next_m_hat = next_m / (1 - beta1_power)
      next_v_hat = next_v / (1 - beta2_power)
      update = next_m_hat / (tf.sqrt(next_v_hat + 1e-10) + self._epsilon)

      next_param = param - lr * update

      assignments.extend(
          [param.assign(next_param),
           m.assign(next_m),
           v.assign(next_v)])

    # Do this after all other assignments are done to prevent a race condition.
    with tf.control_dependencies(assignments):
      assignments.extend([
          beta1_power.assign(beta1_power * self._beta1),
          beta2_power.assign(beta2_power * self._beta2),
          global_step.assign_add(1),
      ])
    return tf.group(*assignments, name=name)
//  [python]
def apply_gradients(self, grads_and_vars, global_step, name=None):
    """Perform an update with the parameters."""

    # we meta-train with 10k steps. When applying to longer problems we want to
    # have a reasonable schedule so we rescale.

    rescale_global_step = float(10000) / self._training_steps * tf.cast(
        global_step, tf.float32)

    beta1_power = tf.get_variable(
        dtype=tf.float32, name="beta1_power", initializer=self._beta1)
    beta2_power = tf.get_variable(
        dtype=tf.float32, name="beta2_power", initializer=self._beta2)

    exp_factor = tf.exp(-self._exponential_decay *
                        tf.cast(rescale_global_step, tf.float32))

    # lr reduction per step.
    linear_factor = tf.maximum(
        1 - self._linear_decay * tf.cast(rescale_global_step, tf.float32), 0.0)

    lr = exp_factor * linear_factor * self._learning_rate

    assignments = []
    for (grad, param) in grads_and_vars:
      if grad is None or param is None:
        continue
      # sparse to dense conversion
      grad = tf.convert_to_tensor(grad)

      param_name = self._get_variable_name(param.name)

      m = tf.get_variable(
          name=param_name + "/adam_m",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())
      v = tf.get_variable(
          name=param_name + "/adam_v",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())

      next_m = (self._beta1 * m + (1.0 - self._beta1) * grad)
      next_v = (self._beta2 * v + (1.0 - self._beta2) * tf.square(grad))
      next_m_hat = next_m / (1 - beta1_power)
      next_v_hat = next_v / (1 - beta2_power)
      update = next_m_hat / (tf.sqrt(next_v_hat + 1e-10) + self._epsilon)

      next_param = param - lr * update

      assignments.extend(
          [param.assign(next_param),
           m.assign(next_m),
           v.assign(next_v)])

    # Do this after all other assignments are done to prevent a race condition.
    with tf.control_dependencies(assignments):
      assignments.extend([
          beta1_power.assign(beta1_power * self._beta1),
          beta2_power.assign(beta2_power * self._beta2),
          global_step.assign_add(1),
      ])
    return tf.group(*assignments, name=name)
//  [python]
def _ignore_pad(embeddings_table, ids, use_one_hot_embeddings=False):
  """Use mean of symbol embeddings as overall embedding but ignore PAD."""
  source_embeddings = common_layers.embedding_lookup(embeddings_table, ids,
                                                     use_one_hot_embeddings)
  # Set weights to ignore padding.
  embedded_weights = tf.to_float(tf.not_equal(ids, constants.PAD_SYMBOL_ID))
  embedded_weights = tf.expand_dims(embedded_weights, -1)
  return source_embeddings * embedded_weights
//  [python]
def _ignore_pad(embeddings_table, ids, use_one_hot_embeddings=False):
  """Use mean of symbol embeddings as overall embedding but ignore PAD."""
  source_embeddings = common_layers.embedding_lookup(embeddings_table, ids,
                                                     use_one_hot_embeddings)
  # Set weights to ignore padding.
  embedded_weights = tf.cast(tf.not_equal(ids, constants.PAD_SYMBOL_ID), tf.float32)
  embedded_weights = tf.expand_dims(embedded_weights, -1)
  return source_embeddings * embedded_weights
//  [python]
def get_image_label_from_record(image_record, label_record):
  """Decodes the image and label information from one data record."""
  # Convert from tf.string to tf.uint8.
  image = tf.decode_raw(image_record, tf.uint8)
  # Convert from tf.uint8 to tf.float32.
  image = tf.cast(image, tf.float32)

  # Reshape image to correct shape.
  image = tf.reshape(image, [IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS])
  # Normalize from [0, 255] to [0.0, 1.0]
  image /= 255.0

  # Convert from tf.string to tf.uint8.
  label = tf.decode_raw(label_record, tf.uint8)
  # Convert from tf.uint8 to tf.int32.
  label = tf.to_int32(label)
  # Reshape label to correct shape.
  label = tf.reshape(label, [])  # label is a scalar
  return image, label
//  [python]
def get_image_label_from_record(image_record, label_record):
  """Decodes the image and label information from one data record."""
  # Convert from tf.string to tf.uint8.
  image = tf.decode_raw(image_record, tf.uint8)
  # Convert from tf.uint8 to tf.float32.
  image = tf.cast(image, tf.float32)

  # Reshape image to correct shape.
  image = tf.reshape(image, [IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS])
  # Normalize from [0, 255] to [0.0, 1.0]
  image /= 255.0

  # Convert from tf.string to tf.uint8.
  label = tf.decode_raw(label_record, tf.uint8)
  # Convert from tf.uint8 to tf.int32.
  label = tf.cast(label, tf.int32)
  # Reshape label to correct shape.
  label = tf.reshape(label, [])  # label is a scalar
  return image, label
//  [python]
def get_timing_signal_1d(
    length,
    channels,
    min_timescale = 1.0,
    max_timescale = 1.0e4,
    start_index = 0,
):
    """Gets a bunch of sinusoids of different frequencies.

  Each channel of the input Tensor is incremented by a sinusoid of a different
  frequency and phase.

  This allows attention to learn to use absolute and relative positions.
  Timing signals should be added to some precursors of both the query and the
  memory inputs to attention.

  The use of relative position is possible because sin(x+y) and cos(x+y) can be
  expressed in terms of y, sin(x) and cos(x).

  In particular, we use a geometric sequence of timescales starting with
  min_timescale and ending with max_timescale.  The number of different
  timescales is equal to channels / 2. For each timescale, we
  generate the two sinusoidal signals sin(timestep/timescale) and
  cos(timestep/timescale).  All of these sinusoids are concatenated in
  the channels dimension.

  Args:
    length: scalar, length of timing signal sequence.
    channels: scalar, size of timing embeddings to create. The number of
      different timescales is equal to channels / 2.
    min_timescale: a float
    max_timescale: a float
    start_index: index of first position

  Returns:
    a Tensor of timing signals [1, length, channels]
  """
    position = tf.to_float(tf.range(length) + start_index)
    num_timescales = channels // 2
    log_timescale_increment = math.log(
        float(max_timescale) / float(min_timescale)
    ) / tf.maximum(tf.to_float(num_timescales) - 1, 1)
    inv_timescales = min_timescale * tf.exp(
        tf.to_float(tf.range(num_timescales)) * -log_timescale_increment
    )
    scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(
        inv_timescales, 0
    )
    # Please note that this slightly differs from the published paper.
    # See a discussion here: https://github.com/tensorflow/tensor2tensor/pull/177
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis = 1)
    signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
    signal = tf.reshape(signal, [1, length, channels])
    return signal
//  [python]
def get_timing_signal_1d(
    length,
    channels,
    min_timescale = 1.0,
    max_timescale = 1.0e4,
    start_index = 0,
):
    """Gets a bunch of sinusoids of different frequencies.

  Each channel of the input Tensor is incremented by a sinusoid of a different
  frequency and phase.

  This allows attention to learn to use absolute and relative positions.
  Timing signals should be added to some precursors of both the query and the
  memory inputs to attention.

  The use of relative position is possible because sin(x+y) and cos(x+y) can be
  expressed in terms of y, sin(x) and cos(x).

  In particular, we use a geometric sequence of timescales starting with
  min_timescale and ending with max_timescale.  The number of different
  timescales is equal to channels / 2. For each timescale, we
  generate the two sinusoidal signals sin(timestep/timescale) and
  cos(timestep/timescale).  All of these sinusoids are concatenated in
  the channels dimension.

  Args:
    length: scalar, length of timing signal sequence.
    channels: scalar, size of timing embeddings to create. The number of
      different timescales is equal to channels / 2.
    min_timescale: a float
    max_timescale: a float
    start_index: index of first position

  Returns:
    a Tensor of timing signals [1, length, channels]
  """
    position = tf.cast(tf.range(length) + start_index, dtype=tf.float32)
    num_timescales = channels // 2
    log_timescale_increment = math.log(
        float(max_timescale) / float(min_timescale)
    ) / tf.maximum(tf.cast(num_timescales - 1, dtype=tf.float32), 1)
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), dtype=tf.float32) * -log_timescale_increment
    )
    scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(
        inv_timescales, 0
    )
    # Please note that this slightly differs from the published paper.
    # See a discussion here: https://github.com/tensorflow/tensor2tensor/pull/177
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis = 1)
    signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
    signal = tf.reshape(signal, [1, length, channels])
    return signal
//  [python]
def _classification_loss(self, pred_label, gt_label, num_matched_boxes):
    """Computes the classification loss.

    Computes the classification loss with hard negative mining.
    Args:
      pred_label: a flatten tensor that includes all predicted class. The shape
        is [batch_size, num_anchors, num_classes].
      gt_label: a tensor that represents the classification groundtruth targets.
        The shape is [batch_size, num_anchors, 1].
      num_matched_boxes: the number of anchors that are matched to a groundtruth
        targets. This is used as the loss normalizater.

    Returns:
      box_loss: a float32 representing total box regression loss.
    """
    cross_entropy = tf.losses.sparse_softmax_cross_entropy(
        gt_label, pred_label, reduction=tf.losses.Reduction.NONE)

    mask = tf.greater(tf.squeeze(gt_label), 0)
    float_mask = tf.cast(mask, tf.float32)

    # Hard example mining
    neg_masked_cross_entropy = cross_entropy * (1 - float_mask)
    relative_position = tf.argsort(
        tf.argsort(
            neg_masked_cross_entropy, direction='DESCENDING'))
    num_neg_boxes = tf.minimum(
        tf.to_int32(num_matched_boxes) * ssd_constants.NEGS_PER_POSITIVE,
        ssd_constants.NUM_SSD_BOXES)
    top_k_neg_mask = tf.cast(tf.less(
        relative_position,
        tf.tile(num_neg_boxes[:, tf.newaxis], (1, ssd_constants.NUM_SSD_BOXES))
    ), tf.float32)

    class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)

    return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def _classification_loss(self, pred_label, gt_label, num_matched_boxes):
    """Computes the classification loss.

    Computes the classification loss with hard negative mining.
    Args:
      pred_label: a flatten tensor that includes all predicted class. The shape
        is [batch_size, num_anchors, num_classes].
      gt_label: a tensor that represents the classification groundtruth targets.
        The shape is [batch_size, num_anchors, 1].
      num_matched_boxes: the number of anchors that are matched to a groundtruth
        targets. This is used as the loss normalizater.

    Returns:
      box_loss: a float32 representing total box regression loss.
    """
    cross_entropy = tf.losses.sparse_softmax_cross_entropy(
        gt_label, pred_label, reduction=tf.losses.Reduction.NONE)

    mask = tf.greater(tf.squeeze(gt_label), 0)
    float_mask = tf.cast(mask, tf.float32)

    # Hard example mining
    neg_masked_cross_entropy = cross_entropy * (1 - float_mask)
    relative_position = tf.argsort(
        tf.argsort(
            neg_masked_cross_entropy, direction='DESCENDING'))
    num_neg_boxes = tf.minimum(
        tf.cast(num_matched_boxes, tf.int32) * ssd_constants.NEGS_PER_POSITIVE,
        ssd_constants.NUM_SSD_BOXES)
    top_k_neg_mask = tf.cast(tf.less(
        relative_position,
        tf.tile(num_neg_boxes[:, tf.newaxis], (1, ssd_constants.NUM_SSD_BOXES))
    ), tf.float32)

    class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)

    return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def _bert_embeddings(wordpiece_embedding_size, bert_config, features,
                     is_training, use_one_hot_embeddings, scope,
                     use_segment_ids):
  """Get embeddings from BERT."""
  token_type_ids = None
  if use_segment_ids:
    token_type_ids = features[constants.SEGMENT_ID_KEY]

  max_seq_len = tf.shape(features[constants.SOURCE_WORDPIECES_KEY])[1]
  input_mask = bert_utils.get_input_mask(max_seq_len,
                                         features[constants.SOURCE_LEN_KEY])
  input_ids = features[constants.SOURCE_WORDPIECES_KEY]
  source_embeddings = bert_utils.get_bert_embeddings(
      input_ids,
      bert_config,
      input_mask,
      token_type_ids=token_type_ids,
      is_training=is_training,
      use_one_hot_embeddings=use_one_hot_embeddings,
      scope=scope)
  source_embeddings = common_layers.linear_transform(source_embeddings,
                                                     wordpiece_embedding_size,
                                                     "bert_transform")

  # Set weights to ignore padding.
  embedded_weights = tf.to_float(
      tf.not_equal(input_ids, constants.PAD_SYMBOL_ID))
  embedded_weights = tf.expand_dims(embedded_weights, -1)
  return source_embeddings * embedded_weights
//  [python]
def _bert_embeddings(wordpiece_embedding_size, bert_config, features,
                     is_training, use_one_hot_embeddings, scope,
                     use_segment_ids):
  """Get embeddings from BERT."""
  token_type_ids = None
  if use_segment_ids:
    token_type_ids = features[constants.SEGMENT_ID_KEY]

  max_seq_len = tf.shape(features[constants.SOURCE_WORDPIECES_KEY])[1]
  input_mask = bert_utils.get_input_mask(max_seq_len,
                                         features[constants.SOURCE_LEN_KEY])
  input_ids = features[constants.SOURCE_WORDPIECES_KEY]
  source_embeddings = bert_utils.get_bert_embeddings(
      input_ids,
      bert_config,
      input_mask,
      token_type_ids=token_type_ids,
      is_training=is_training,
      use_one_hot_embeddings=use_one_hot_embeddings,
      scope=scope)
  source_embeddings = common_layers.linear_transform(source_embeddings,
                                                     wordpiece_embedding_size,
                                                     "bert_transform")

  # Set weights to ignore padding.
  embedded_weights = tf.cast(
      tf.not_equal(input_ids, constants.PAD_SYMBOL_ID), tf.float32)
  embedded_weights = tf.expand_dims(embedded_weights, -1)
  return source_embeddings * embedded_weights
//  [python]
def blend(image1, image2, factor):
    """Blend image1 and image2 using 'factor'.

    Factor can be above 0.0.  A value of 0.0 means only image1 is used.
    A value of 1.0 means only image2 is used.  A value between 0.0 and
    1.0 means we linearly interpolate the pixel values between the two
    images.  A value greater than 1.0 "extrapolates" the difference
    between the two pixel values, and we clip the results to values
    between 0 and 255.

    Args:
        image1: An image Tensor of type uint8.
        image2: An image Tensor of type uint8.
        factor: A floating point value above 0.0.

    Returns:
        A blended image Tensor of type uint8.
    """
    # if factor == 0.0:
    #     return tf.convert_to_tensor(image1)
    # if factor == 1.0:
    #     return tf.convert_to_tensor(image2)
    image1 = tf.to_float(image1)
    image2 = tf.to_float(image2)

    difference = image2 - image1
    scaled = factor * difference

    # Do addition in float.
    temp = tf.to_float(image1) + scaled
    # # Interpolate
    # if 1.0 > factor > 0.0:
    #     # Interpolation means we always stay within 0 and 255.
    #     return tf.cast(temp, tf.uint8)

    # # Extrapolate:
    # # We need to clip and then cast.
    # return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
    result = tf.cond(tf.math.logical_and(tf.greater(factor, 0.0), tf.less(factor, 1.0)),
                     lambda: tf.cast(temp, tf.uint8),
                     lambda: tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8))
    return result
//  [python]
def blend(image1, image2, factor):
    """Blend image1 and image2 using 'factor'.

    Factor can be above 0.0.  A value of 0.0 means only image1 is used.
    A value of 1.0 means only image2 is used.  A value between 0.0 and
    1.0 means we linearly interpolate the pixel values between the two
    images.  A value greater than 1.0 "extrapolates" the difference
    between the two pixel values, and we clip the results to values
    between 0 and 255.

    Args:
        image1: An image Tensor of type uint8.
        image2: An image Tensor of type uint8.
        factor: A floating point value above 0.0.

    Returns:
        A blended image Tensor of type uint8.
    """
    image1 = tf.cast(image1, tf.float32)
    image2 = tf.cast(image2, tf.float32)

    difference = image2 - image1
    scaled = factor * difference

    temp = tf.cast(image1, tf.float32) + scaled
    result = tf.cond(tf.math.logical_and(tf.greater(factor, 0.0), tf.less(factor, 1.0)),
                     lambda: tf.cast(temp, tf.uint8),
                     lambda: tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8))
    return result
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.

  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.
  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.

  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.to_float(image1)
  image2 = tf.to_float(image2)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.to_float(image1) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.

  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.
  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.

  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.cast(image1, tf.float32)
  image2 = tf.cast(image2, tf.float32)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.cast(image1, tf.float32) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def _apply_bbox_augmentation(image, bbox, augmentation_func, *args):
    """Applies augmentation_func to the subsection of image indicated by bbox.

    Args:
        image: 3D uint8 Tensor.
        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
            of type float that represents the normalized coordinates between 0 and 1.
        augmentation_func: Augmentation function that will be applied to the
            subsection of image.
        *args: Additional parameters that will be passed into augmentation_func
            when it is called.

    Returns:
        A modified version of image, where the bbox location in the image will
        have `augmentation_func applied to it.
    """
    image_height = tf.to_float(tf.shape(image)[0])
    image_width = tf.to_float(tf.shape(image)[1])
    min_y = tf.to_int32(image_height * bbox[0])
    min_x = tf.to_int32(image_width * bbox[1])
    max_y = tf.to_int32(image_height * bbox[2])
    max_x = tf.to_int32(image_width * bbox[3])
    image_height = tf.to_int32(image_height)
    image_width = tf.to_int32(image_width)

    # Clip to be sure the max values do not fall out of range.
    max_y = tf.minimum(max_y, image_height - 1)
    max_x = tf.minimum(max_x, image_width - 1)

    # Get the sub-tensor that is the image within the bounding box region.
    bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]

    # Apply the augmentation function to the bbox portion of the image.
    augmented_bbox_content = augmentation_func(bbox_content, *args)

    # Pad the augmented_bbox_content and the mask to match the shape of original
    # image.
    augmented_bbox_content = tf.pad(
        augmented_bbox_content,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]])

    # Create a mask that will be used to zero out a part of the original image.
    mask_tensor = tf.zeros_like(bbox_content)

    mask_tensor = tf.pad(
        mask_tensor,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]],
        constant_values=1)
    # Replace the old bbox content with the new augmented content.
    image = image * mask_tensor + augmented_bbox_content
    return image