jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def flop_coeff(op):
  """Computes the coefficient of number of flops associated with a convolution.

  The FLOPs cost of a convolution is given by C * output_depth * input_depth,
  where C = 2 * output_width * output_height * filter_size. The 2 is because we
  have one multiplication and one addition for each convolution weight and
  pixel. This function returns C.

  Supported operations names are listed in cost_calculator.FLOP_OPS.

  Args:
    op: A tf.Operation of supported types.

  Returns:
    A float, the coefficient that when multiplied by the input depth and by the
    output depth gives the number of flops needed to compute the convolution.

  Raises:
    ValueError: conv_op is not a supported tf.Operation.
  """
  if not is_flop_op(op):
    return 0.0
  if op.type == 'MatMul':
    # A MatMul is like a 1x1 conv with an output size of 1x1, so from the factor
    # below only the 2.0 remains.
    return 2.0
  # Looking at the output shape makes it easy to automatically take into
  # account strides and the type of padding.
  def kernel_num_elements(tensor):
    """Returns the number of elements of a kernel.

    Args:
      tensor: The weight tensor.

    Returns:
      Number of elements of the kernel (either float or tf.float).
    """
    num_elements = np.prod(tensor.shape.dims[1:-1]).value
    if num_elements:
      return num_elements
    return tf.to_float(tf.reduce_prod(tf.shape(tensor)[1:-1]))

  if op.type in ('Conv2D', 'DepthwiseConv2dNative', 'Conv3D'):
    num_elements = kernel_num_elements(op.outputs[0])
  elif op.type == 'Conv2DBackpropInput':
    # For a transposed convolution, the input and the output are swapped (as
    # far as shapes are concerned). In other words, for a given filter shape
    # and stride, if Conv2D maps from shapeX to shapeY, Conv2DBackpropInput
    # maps from shapeY to shapeX. Therefore wherever we use the output shape
    # for Conv2D, we use the input shape for Conv2DBackpropInput.
    num_elements = kernel_num_elements(cost_calculator.get_input_activation(op))
  else:
    # Can only happen if elements are added to FLOP_OPS and not taken care of.
    assert False, '%s in cost_calculator.FLOP_OPS but not handled' % op.type
  # Handle dynamic shaping while keeping old code path to not break
  # other clients.
  return 2.0 * num_elements * _get_conv_filter_size(op)
//  [python]
def flop_coeff(op):
  """Computes the coefficient of number of flops associated with a convolution.

  The FLOPs cost of a convolution is given by C * output_depth * input_depth,
  where C = 2 * output_width * output_height * filter_size. The 2 is because we
  have one multiplication and one addition for each convolution weight and
  pixel. This function returns C.

  Supported operations names are listed in cost_calculator.FLOP_OPS.

  Args:
    op: A tf.Operation of supported types.

  Returns:
    A float, the coefficient that when multiplied by the input depth and by the
    output depth gives the number of flops needed to compute the convolution.

  Raises:
    ValueError: conv_op is not a supported tf.Operation.
  """
  if not is_flop_op(op):
    return 0.0
  if op.type == 'MatMul':
    # A MatMul is like a 1x1 conv with an output size of 1x1, so from the factor
    # below only the 2.0 remains.
    return 2.0
  # Looking at the output shape makes it easy to automatically take into
  # account strides and the type of padding.
  def kernel_num_elements(tensor):
    """Returns the number of elements of a kernel.

    Args:
      tensor: The weight tensor.

    Returns:
      Number of elements of the kernel (either float or tf.float).
    """
    num_elements = np.prod(tensor.shape.dims[1:-1]).value
    if num_elements:
      return num_elements
    return tf.cast(tf.reduce_prod(tf.shape(tensor)[1:-1]), tf.float32)

  if op.type in ('Conv2D', 'DepthwiseConv2dNative', 'Conv3D'):
    num_elements = kernel_num_elements(op.outputs[0])
  elif op.type == 'Conv2DBackpropInput':
    # For a transposed convolution, the input and the output are swapped (as
    # far as shapes are concerned). In other words, for a given filter shape
    # and stride, if Conv2D maps from shapeX to shapeY, Conv2DBackpropInput
    # maps from shapeY to shapeX. Therefore wherever we use the output shape
    # for Conv2D, we use the input shape for Conv2DBackpropInput.
    num_elements = kernel_num_elements(cost_calculator.get_input_activation(op))
  else:
    # Can only happen if elements are added to FLOP_OPS and not taken care of.
    assert False, '%s in cost_calculator.FLOP_OPS but not handled' % op.type
  # Handle dynamic shaping while keeping old code path to not break
  # other clients.
  return 2.0 * num_elements * _get_conv_filter_size(op)
//  [python]
def sequence_loss_fn(self, chars_logits, chars_labels):
    """Loss function for char sequence.

    Depending on values of hyper parameters it applies label smoothing and can
    also ignore all null chars after the first one.

    Args:
      chars_logits: logits for predicted characters,
        shape=[batch_size, seq_length, num_char_classes];
      chars_labels: ground truth ids of characters,
        shape=[batch_size, seq_length];
      mparams: method hyper parameters.

    Returns:
      A Tensor with shape [batch_size] - the log-perplexity for each sequence.
    """
    mparams = self._mparams['sequence_loss_fn']
    with tf.variable_scope('sequence_loss_fn/SLF'):
      if mparams.label_smoothing > 0:
        smoothed_one_hot_labels = self.label_smoothing_regularization(
          chars_labels, mparams.label_smoothing)
        labels_list = tf.unstack(smoothed_one_hot_labels, axis=1)
      else:
        # NOTE: in case of sparse softmax we are not using one-hot
        # encoding.
        labels_list = tf.unstack(chars_labels, axis=1)

      batch_size, seq_length, _ = chars_logits.shape.as_list()
      if mparams.ignore_nulls:
        weights = tf.ones((batch_size, seq_length), dtype=tf.float32)
      else:
        # Suppose that reject character is the last in the charset.
        reject_char = tf.constant(
          self._params.num_char_classes - 1,
          shape=(batch_size, seq_length),
          dtype=tf.int64)
        known_char = tf.not_equal(chars_labels, reject_char)
        weights = tf.to_float(known_char)

      logits_list = tf.unstack(chars_logits, axis=1)
      weights_list = tf.unstack(weights, axis=1)
      loss = tf.contrib.legacy_seq2seq.sequence_loss(
        logits_list,
        labels_list,
        weights_list,
        softmax_loss_function=get_softmax_loss_fn(mparams.label_smoothing),
        average_across_timesteps=mparams.average_across_timesteps)
      tf.losses.add_loss(loss)
      return loss
//  [python]
def sequence_loss_fn(self, chars_logits, chars_labels):
    """Loss function for char sequence.

    Depending on values of hyper parameters it applies label smoothing and can
    also ignore all null chars after the first one.

    Args:
      chars_logits: logits for predicted characters,
        shape=[batch_size, seq_length, num_char_classes];
      chars_labels: ground truth ids of characters,
        shape=[batch_size, seq_length];
      mparams: method hyper parameters.

    Returns:
      A Tensor with shape [batch_size] - the log-perplexity for each sequence.
    """
    mparams = self._mparams['sequence_loss_fn']
    with tf.variable_scope('sequence_loss_fn/SLF'):
      if mparams.label_smoothing > 0:
        smoothed_one_hot_labels = self.label_smoothing_regularization(
          chars_labels, mparams.label_smoothing)
        labels_list = tf.unstack(smoothed_one_hot_labels, axis=1)
      else:
        # NOTE: in case of sparse softmax we are not using one-hot
        # encoding.
        labels_list = tf.unstack(chars_labels, axis=1)

      batch_size, seq_length, _ = chars_logits.shape.as_list()
      if mparams.ignore_nulls:
        weights = tf.ones((batch_size, seq_length), dtype=tf.float32)
      else:
        # Suppose that reject character is the last in the charset.
        reject_char = tf.constant(
          self._params.num_char_classes - 1,
          shape=(batch_size, seq_length),
          dtype=tf.int64)
        known_char = tf.not_equal(chars_labels, reject_char)
        weights = tf.cast(known_char, tf.float32)

      logits_list = tf.unstack(chars_logits, axis=1)
      weights_list = tf.unstack(weights, axis=1)
      loss = tf.contrib.legacy_seq2seq.sequence_loss(
        logits_list,
        labels_list,
        weights_list,
        softmax_loss_function=get_softmax_loss_fn(mparams.label_smoothing),
        average_across_timesteps=mparams.average_across_timesteps)
      tf.losses.add_loss(loss)
      return loss
//  [python]
def softmax_loss(self, antecedent_scores, antecedent_labels):
        gold_scores = antecedent_scores + \
            tf.log(tf.to_float(antecedent_labels))  # [k, max_ant + 1]
        marginalized_gold_scores = tf.reduce_logsumexp(gold_scores, [1])  # [k]
        log_norm = tf.reduce_logsumexp(antecedent_scores, [1])  # [k]
        return log_norm - marginalized_gold_scores
//  [python]
def softmax_loss(self, antecedent_scores, antecedent_labels):
        gold_scores = antecedent_scores + \
            tf.log(tf.cast(antecedent_labels, dtype=tf.float32))  # [k, max_ant + 1]
        marginalized_gold_scores = tf.reduce_logsumexp(gold_scores, [1])  # [k]
        log_norm = tf.reduce_logsumexp(antecedent_scores, [1])  # [k]
        return log_norm - marginalized_gold_scores
//  [python]
def distance_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_antecedent_offsets = tf.tile(
            tf.expand_dims(tf.range(c) + 1, 0), [k, 1])  # [k, c]
        raw_top_antecedents = tf.expand_dims(
            tf.range(k), 1) - top_antecedent_offsets  # [k, c]
        top_antecedents_mask = raw_top_antecedents >= 0  # [k, c]
        top_antecedents = tf.maximum(raw_top_antecedents, 0)  # [k, c]

        # [k, c]
        top_fast_antecedent_scores = tf.expand_dims(
            top_span_mention_scores, 1) + tf.gather(top_span_mention_scores,
                                                    top_antecedents)
        # [k, c]
        top_fast_antecedent_scores += tf.log(tf.to_float(top_antecedents_mask))
        return (top_antecedents, top_antecedents_mask,
                top_fast_antecedent_scores, top_antecedent_offsets)
//  [python]
def distance_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_antecedent_offsets = tf.tile(
            tf.expand_dims(tf.range(c) + 1, 0), [k, 1])  # [k, c]
        raw_top_antecedents = tf.expand_dims(
            tf.range(k), 1) - top_antecedent_offsets  # [k, c]
        top_antecedents_mask = raw_top_antecedents >= 0  # [k, c]
        top_antecedents = tf.maximum(raw_top_antecedents, 0)  # [k, c]

        # [k, c]
        top_fast_antecedent_scores = tf.expand_dims(
            top_span_mention_scores, 1) + tf.gather(top_span_mention_scores,
                                                    top_antecedents)
        # [k, c]
        top_fast_antecedent_scores += tf.log(tf.cast(top_antecedents_mask, dtype=tf.float32))
        return (top_antecedents, top_antecedents_mask,
                top_fast_antecedent_scores, top_antecedent_offsets)
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into
        10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.to_int32(
            tf.floor(tf.log(tf.to_float(distances)) / math.log(2))) + 3
        use_identity = tf.to_int32(distances <= 4)
        combined_idx = use_identity * distances + \
            (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into
        10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.cast(
            tf.floor(tf.math.log(tf.cast(distances, tf.float32)) / math.log(2)), tf.int32) + 3
        use_identity = tf.cast(distances <= 4, tf.int32)
        combined_idx = use_identity * distances + \
            (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def pad_to_multiple(tensor, multiple):
  """Returns the tensor zero padded to the specified multiple.

  Appends 0s to the end of the first and second dimension (height and width) of
  the tensor until both dimensions are a multiple of the input argument
  'multiple'. E.g. given an input tensor of shape [1, 3, 5, 1] and an input
  multiple of 4, PadToMultiple will append 0s so that the resulting tensor will
  be of shape [1, 4, 8, 1].

  Args:
    tensor: rank 4 float32 tensor, where
            tensor -> [batch_size, height, width, channels].
    multiple: the multiple to pad to.

  Returns:
    padded_tensor: the tensor zero padded to the specified multiple.
  """
  tensor_shape = tensor.get_shape()
  batch_size = static_shape.get_batch_size(tensor_shape)
  tensor_height = static_shape.get_height(tensor_shape)
  tensor_width = static_shape.get_width(tensor_shape)
  tensor_depth = static_shape.get_depth(tensor_shape)

  if batch_size is None:
    batch_size = tf.shape(tensor)[0]

  if tensor_height is None:
    tensor_height = tf.shape(tensor)[1]
    padded_tensor_height = tf.to_int32(
        tf.ceil(tf.to_float(tensor_height) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_height = int(
        math.ceil(float(tensor_height) / multiple) * multiple)

  if tensor_width is None:
    tensor_width = tf.shape(tensor)[2]
    padded_tensor_width = tf.to_int32(
        tf.ceil(tf.to_float(tensor_width) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_width = int(
        math.ceil(float(tensor_width) / multiple) * multiple)

  if tensor_depth is None:
    tensor_depth = tf.shape(tensor)[3]

  # Use tf.concat instead of tf.pad to preserve static shape
  height_pad = tf.zeros([
      batch_size, padded_tensor_height - tensor_height, tensor_width,
      tensor_depth
  ])
  padded_tensor = tf.concat([tensor, height_pad], 1)
  width_pad = tf.zeros([
      batch_size, padded_tensor_height, padded_tensor_width - tensor_width,
      tensor_depth
  ])
  padded_tensor = tf.concat([padded_tensor, width_pad], 2)

  return padded_tensor
//  [python]
def pad_to_multiple(tensor, multiple):
  """Returns the tensor zero padded to the specified multiple.

  Appends 0s to the end of the first and second dimension (height and width) of
  the tensor until both dimensions are a multiple of the input argument
  'multiple'. E.g. given an input tensor of shape [1, 3, 5, 1] and an input
  multiple of 4, PadToMultiple will append 0s so that the resulting tensor will
  be of shape [1, 4, 8, 1].

  Args:
    tensor: rank 4 float32 tensor, where
            tensor -> [batch_size, height, width, channels].
    multiple: the multiple to pad to.

  Returns:
    padded_tensor: the tensor zero padded to the specified multiple.
  """
  tensor_shape = tensor.get_shape()
  batch_size = static_shape.get_batch_size(tensor_shape)
  tensor_height = static_shape.get_height(tensor_shape)
  tensor_width = static_shape.get_width(tensor_shape)
  tensor_depth = static_shape.get_depth(tensor_shape)

  if batch_size is None:
    batch_size = tf.shape(tensor)[0]

  if tensor_height is None:
    tensor_height = tf.shape(tensor)[1]
    padded_tensor_height = tf.cast(
        tf.ceil(tf.to_float(tensor_height) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_height = int(
        math.ceil(float(tensor_height) / multiple) * multiple)

  if tensor_width is None:
    tensor_width = tf.shape(tensor)[2]
    padded_tensor_width = tf.cast(
        tf.ceil(tf.to_float(tensor_width) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_width = int(
        math.ceil(float(tensor_width) / multiple) * multiple)

  if tensor_depth is None:
    tensor_depth = tf.shape(tensor)[3]

  # Use tf.concat instead of tf.pad to preserve static shape
  height_pad = tf.zeros([
      batch_size, padded_tensor_height - tensor_height, tensor_width,
      tensor_depth
  ])
  padded_tensor = tf.concat([tensor, height_pad], 1)
  width_pad = tf.zeros([
      batch_size, padded_tensor_height, padded_tensor_width - tensor_width,
      tensor_depth
  ])
  padded_tensor = tf.concat([padded_tensor, width_pad], 2)

  return padded_tensor
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_span_range = tf.range(k)  # [k]
        antecedent_offsets = tf.expand_dims(
            top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
        antecedents_mask = antecedent_offsets >= 1  # [k, k]
        # [k, k]
        fast_antecedent_scores = (tf.expand_dims(top_span_mention_scores, 1)
                                  + tf.expand_dims(top_span_mention_scores, 0))
        # [k, k]
        fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask))
        # [k, k]
        fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)

        _, top_antecedents = tf.nn.top_k(
            fast_antecedent_scores, c, sorted=False)  # [k, c]
        top_antecedents_mask = util.batch_gather(
            antecedents_mask, top_antecedents)  # [k, c]
        top_fast_antecedent_scores = util.batch_gather(
            fast_antecedent_scores, top_antecedents)  # [k, c]
        top_antecedent_offsets = util.batch_gather(
            antecedent_offsets, top_antecedents)  # [k, c]
        return (top_antecedents, top_antecedents_mask,
                top_fast_antecedent_scores, top_antecedent_offsets)
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_span_range = tf.range(k)  # [k]
        antecedent_offsets = tf.expand_dims(
            top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
        antecedents_mask = antecedent_offsets >= 1  # [k, k]
        # [k, k]
        fast_antecedent_scores = (tf.expand_dims(top_span_mention_scores, 1)
                                  + tf.expand_dims(top_span_mention_scores, 0))
        # [k, k]
        fast_antecedent_scores += tf.log(tf.cast(antecedents_mask, dtype=tf.float32))
        # [k, k]
        fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)

        _, top_antecedents = tf.nn.top_k(
            fast_antecedent_scores, c, sorted=False)  # [k, c]
        top_antecedents_mask = util.batch_gather(
            antecedents_mask, top_antecedents)  # [k, c]
        top_fast_antecedent_scores = util.batch_gather(
            fast_antecedent_scores, top_antecedents)  # [k, c]
        top_antecedent_offsets = util.batch_gather(
            antecedent_offsets, top_antecedents)  # [k, c]
        return (top_antecedents, top_antecedents_mask,
                top_fast_antecedent_scores, top_antecedent_offsets)
//  [python]
def load_bernoulli_mnist_dataset(directory, split_name):
  """Returns Hugo Larochelle's binary static MNIST tf.data.Dataset."""
  amat_file = download(directory, FILE_TEMPLATE.format(split=split_name))
  dataset = tf.data.TextLineDataset(amat_file)
  str_to_arr = lambda string: np.array([c == b"1" for c in string.split()])

  def _parser(s):
    booltensor = tf.py_func(str_to_arr, [s], tf.bool)
    reshaped = tf.reshape(booltensor, [28, 28, 1])
    return tf.to_float(reshaped), tf.constant(0, tf.int32)

  return dataset.map(_parser)
//  [python]
def load_bernoulli_mnist_dataset(directory, split_name):
  """Returns Hugo Larochelle's binary static MNIST tf.data.Dataset."""
  amat_file = download(directory, FILE_TEMPLATE.format(split=split_name))
  dataset = tf.data.TextLineDataset(amat_file)
  str_to_arr = lambda string: np.array([c == b"1" for c in string.split()])

  def _parser(s):
    booltensor = tf.py_func(str_to_arr, [s], tf.bool)
    reshaped = tf.reshape(booltensor, [28, 28, 1])
    return tf.cast(reshaped, tf.float32), tf.constant(0, tf.int32)

  return dataset.map(_parser)
//  [python]
def preprocess_example(self, example, mode, _):
    # Resize from usual size ~1350x60 to 90x4 in this test.
    img = example["inputs"]
    img = tf.to_int64(
        tf.image.resize_images(img, [90, 4], tf.image.ResizeMethod.AREA))
    img = tf.image.per_image_standardization(img)
    example["inputs"] = img
    return example
//  [python]
def preprocess_example(self, example, mode, _):
    # Resize from usual size ~1350x60 to 90x4 in this test.
    img = example["inputs"]
    img = tf.cast(
        tf.image.resize_images(img, [90, 4], tf.image.ResizeMethod.AREA), tf.int64)
    img = tf.image.per_image_standardization(img)
    example["inputs"] = img
    return example
//  [python]
def bpr_mf(user_count,item_count,hidden_dim):
    u = tf.placeholder(tf.int32,[None])
    i = tf.placeholder(tf.int32,[None])
    j = tf.placeholder(tf.int32,[None])

    user_emb_w = tf.get_variable("user_emb_w", [user_count + 1, hidden_dim],
                                 initializer=tf.random_normal_initializer(0, 0.1))
    item_emb_w = tf.get_variable("item_emb_w", [item_count + 1, hidden_dim],
                                 initializer=tf.random_normal_initializer(0, 0.1))

    u_emb = tf.nn.embedding_lookup(user_emb_w, u)
    i_emb = tf.nn.embedding_lookup(item_emb_w, i)
    j_emb = tf.nn.embedding_lookup(item_emb_w, j)


    x = tf.reduce_sum(tf.multiply(u_emb,(i_emb-j_emb)),1,keep_dims=True)

    mf_auc = tf.reduce_mean(tf.to_float(x>0))

    l2_norm = tf.add_n([
        tf.reduce_sum(tf.multiply(u_emb, u_emb)),
        tf.reduce_sum(tf.multiply(i_emb, i_emb)),
        tf.reduce_sum(tf.multiply(j_emb, j_emb))
    ])

    regulation_rate = 0.0001
    bprloss = regulation_rate * l2_norm - tf.reduce_mean(tf.log(tf.sigmoid(x)))

    train_op = tf.train.GradientDescentOptimizer(0.01).minimize(bprloss)
    return u, i, j, mf_auc, bprloss, train_op
//  [python]
def bpr_mf(user_count,item_count,hidden_dim):
    u = tf.placeholder(tf.int32,[None])
    i = tf.placeholder(tf.int32,[None])
    j = tf.placeholder(tf.int32,[None])

    user_emb_w = tf.get_variable("user_emb_w", [user_count + 1, hidden_dim],
                                 initializer=tf.random_normal_initializer(0, 0.1))
    item_emb_w = tf.get_variable("item_emb_w", [item_count + 1, hidden_dim],
                                 initializer=tf.random_normal_initializer(0, 0.1))

    u_emb = tf.nn.embedding_lookup(user_emb_w, u)
    i_emb = tf.nn.embedding_lookup(item_emb_w, i)
    j_emb = tf.nn.embedding_lookup(item_emb_w, j)


    x = tf.reduce_sum(tf.multiply(u_emb,(i_emb-j_emb)),1,keep_dims=True)

    mf_auc = tf.reduce_mean(tf.cast(x>0, tf.float32))

    l2_norm = tf.add_n([
        tf.reduce_sum(tf.multiply(u_emb, u_emb)),
        tf.reduce_sum(tf.multiply(i_emb, i_emb)),
        tf.reduce_sum(tf.multiply(j_emb, j_emb))
    ])

    regulation_rate = 0.0001
    bprloss = regulation_rate * l2_norm - tf.reduce_mean(tf.log(tf.sigmoid(x)))

    train_op = tf.train.GradientDescentOptimizer(0.01).minimize(bprloss)
    return u, i, j, mf_auc, bprloss, train_op
//  [python]
def file_based_input_fn_builder(input_files, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_files)
    if is_training:
        d = d.repeat()
        d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_files, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_files)
    if is_training:
        d = d.repeat()
        d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def __call__(self, class_outputs, class_targets):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is softmax on all RoIs.
    Reference:
    https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py
    # pylint: disable=line-too-long

    Args:
      class_outputs: a float tensor representing the class prediction for each
        box with a shape of [batch_size, num_boxes, num_classes].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].

    Returns:
      a scalar tensor representing total class loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      _, _, num_classes = class_outputs.get_shape().as_list()
      class_targets = tf.to_int32(class_targets)
      class_targets_one_hot = tf.one_hot(class_targets, num_classes)
      return self._fast_rcnn_class_loss(class_outputs, class_targets_one_hot)
//  [python]
def __call__(self, class_outputs, class_targets):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is softmax on all RoIs.
    Reference:
    https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py
    # pylint: disable=line-too-long

    Args:
      class_outputs: a float tensor representing the class prediction for each
        box with a shape of [batch_size, num_boxes, num_classes].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].

    Returns:
      a scalar tensor representing total class loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      _, _, num_classes = class_outputs.get_shape().as_list()
      class_targets = tf.cast(class_targets, tf.int32)
      class_targets_one_hot = tf.one_hot(class_targets, num_classes)
      return self._fast_rcnn_class_loss(class_outputs, class_targets_one_hot)
//  [python]
def softmax_loss(self, cluster_scores, gold_labels):
    marginalized_gold_scores = tf.reduce_logsumexp(cluster_scores + tf.log(tf.to_float(gold_labels)), [1])
    log_norm = tf.reduce_logsumexp(cluster_scores, [1])
    loss = log_norm - marginalized_gold_scores
    return loss
//  [python]
def softmax_loss(self, cluster_scores, gold_labels):
    marginalized_gold_scores = tf.reduce_logsumexp(cluster_scores + tf.log(tf.cast(gold_labels, tf.float32)), [1])
    log_norm = tf.reduce_logsumexp(cluster_scores, [1])
    loss = log_norm - marginalized_gold_scores
    return loss
//  [python]
def bucket_distance(self, distances):
    """
    Places the given values (designed for distances) into 10 semi-logscale buckets:
    [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
    """
    logspace_idx = tf.to_int32(tf.floor(tf.log(tf.to_float(distances)) / math.log(2))) + 3
    use_identity = tf.to_int32(distances <= 4)
    combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
    return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def bucket_distance(self, distances):
    """
    Places the given values (designed for distances) into 10 semi-logscale buckets:
    [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
    """
    logspace_idx = tf.cast(tf.floor(tf.log(tf.cast(distances, tf.float32)) / math.log(2)), tf.int32) + 3
    use_identity = tf.cast(distances <= 4, tf.int32)
    combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
    return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def file_based_input_fn_builder_single(input_file, seq_length, num_candidates,
                                       is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "q_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "cand_nums":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "input_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "input_mask":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "segment_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "label_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder_single(input_file, seq_length, num_candidates,
                                       is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "q_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "cand_nums":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "input_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "input_mask":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "segment_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "label_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def summarize_data_utilization(v, tf_global_step, batch_size, epsilon=0.001):
  """Summarizes the samples of non-zero weights during training.

  Args:
    v: a tensor [batch_size, 1] represents the sample weights.
      0: loss, 1: loss difference to the moving average, 2: label and 3: epoch,
      where epoch is an integer between 0 and 99 (the first and the last epoch).
    tf_global_step: the tensor of the current global step.
    batch_size: an integer batch_size
    epsilon: the rounding error. If the weight is smaller than epsilon then set
      it to zero.
  Returns:
    data_util: a tensor of data utilization.
  """
  nonzero_v = tf.get_variable('data_util/nonzero_v', [],
                              initializer=tf.zeros_initializer(),
                              trainable=False,
                              dtype=tf.float32)

  rounded_v = tf.maximum(v - epsilon, tf.to_float(0))

  # Log data utilization
  nonzero_v = tf.assign_add(nonzero_v, tf.count_nonzero(
      rounded_v, dtype=tf.float32))

  data_util = (nonzero_v) / tf.to_float(batch_size) / (
      tf.to_float(tf_global_step) + 2)
  data_util = tf.minimum(data_util, 1)
  tf.stop_gradient(data_util)

  slim.summaries.add_scalar_summary(data_util, 'data_util/data_util')
  slim.summaries.add_scalar_summary(tf.reduce_sum(v), 'data_util/batch_sum_v')
  return data_util
//  [python]
def summarize_data_utilization(v, tf_global_step, batch_size, epsilon=0.001):
  """Summarizes the samples of non-zero weights during training.

  Args:
    v: a tensor [batch_size, 1] represents the sample weights.
      0: loss, 1: loss difference to the moving average, 2: label and 3: epoch,
      where epoch is an integer between 0 and 99 (the first and the last epoch).
    tf_global_step: the tensor of the current global step.
    batch_size: an integer batch_size
    epsilon: the rounding error. If the weight is smaller than epsilon then set
      it to zero.
  Returns:
    data_util: a tensor of data utilization.
  """
  nonzero_v = tf.get_variable('data_util/nonzero_v', [],
                              initializer=tf.zeros_initializer(),
                              trainable=False,
                              dtype=tf.float32)

  rounded_v = tf.maximum(v - epsilon, tf.cast(0, tf.float32))

  # Log data utilization
  nonzero_v = tf.assign_add(nonzero_v, tf.count_nonzero(
      rounded_v, dtype=tf.float32))

  data_util = (nonzero_v) / tf.cast(batch_size, tf.float32) / (
      tf.cast(tf_global_step, tf.float32) + 2)
  data_util = tf.minimum(data_util, 1)
  tf.stop_gradient(data_util)

  slim.summaries.add_scalar_summary(data_util, 'data_util/data_util')
  slim.summaries.add_scalar_summary(tf.reduce_sum(v), 'data_util/batch_sum_v')
  return data_util
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
    same_start = tf.equal(tf.expand_dims(labeled_starts, 1),
                          tf.expand_dims(candidate_starts, 0))  # [num_labeled, num_candidates]
    same_end = tf.equal(tf.expand_dims(labeled_ends, 1),
                        tf.expand_dims(candidate_ends, 0))  # [num_labeled, num_candidates]
    same_span = tf.logical_and(same_start, same_end)  # [num_labeled, num_candidates]
    candidate_labels = tf.matmul(tf.expand_dims(labels, 0), tf.to_int32(same_span))  # [1, num_candidates]
    candidate_labels = tf.squeeze(candidate_labels, 0)  # [num_candidates]
    return candidate_labels
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
    same_start = tf.equal(tf.expand_dims(labeled_starts, 1),
                          tf.expand_dims(candidate_starts, 0))  # [num_labeled, num_candidates]
    same_end = tf.equal(tf.expand_dims(labeled_ends, 1),
                        tf.expand_dims(candidate_ends, 0))  # [num_labeled, num_candidates]
    same_span = tf.logical_and(same_start, same_end)  # [num_labeled, num_candidates]
    candidate_labels = tf.matmul(tf.expand_dims(labels, 0), tf.cast(same_span, tf.int32))  # [1, num_candidates]
    candidate_labels = tf.squeeze(candidate_labels, 0)  # [num_candidates]
    return candidate_labels
//  [python]
def file_based_input_fn_builder_single_files(input_files,
                                             seq_length,
                                             num_candidates,
                                             is_training,
                                             drop_remainder,
                                             num_cpu_threads=4):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "q_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "cand_nums":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "input_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "input_mask":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "segment_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "label_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "scores":
          tf.FixedLenFeature([num_candidates], tf.float32),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    if is_training:
      d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files))
      d = d.repeat()
      d = d.shuffle(buffer_size=len(input_files))
      cycle_length = min(num_cpu_threads, len(input_files))
      # `sloppy` mode means that the interleaving is not exact. This adds
      # even more randomness to the training pipeline.
      d = d.apply(
          tf.data.experimental.parallel_interleave(
              tf.data.TFRecordDataset,
              sloppy=is_training,
              cycle_length=cycle_length))
      d = d.shuffle(buffer_size=100)

    else:
      d = tf.data.TFRecordDataset(input_files)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            num_parallel_batches=num_cpu_threads,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder_single_files(input_files,
                                             seq_length,
                                             num_candidates,
                                             is_training,
                                             drop_remainder,
                                             num_cpu_threads=4):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "q_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "cand_nums":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "input_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "input_mask":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "segment_ids":
          tf.FixedLenFeature([seq_length * num_candidates], tf.int64),
      "label_ids":
          tf.FixedLenFeature([num_candidates], tf.int64),
      "scores":
          tf.FixedLenFeature([num_candidates], tf.float32),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    if is_training:
      d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files))
      d = d.repeat()
      d = d.shuffle(buffer_size=len(input_files))
      cycle_length = min(num_cpu_threads, len(input_files))
      # `sloppy` mode means that the interleaving is not exact. This adds
      # even more randomness to the training pipeline.
      d = d.apply(
          tf.data.experimental.parallel_interleave(
              tf.data.TFRecordDataset,
              sloppy=is_training,
              cycle_length=cycle_length))
      d = d.shuffle(buffer_size=100)

    else:
      d = tf.data.TFRecordDataset(input_files)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            num_parallel_batches=num_cpu_threads,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def smoothing_cross_entropy(self,logits, labels, vocab_size, confidence=0.9): #confidence = 1.0 - label_smoothing. where label_smooth=0.1. from http://github.com/tensorflow/tensor2tensor
        """Cross entropy with label smoothing to limit over-confidence."""
        with tf.name_scope("smoothing_cross_entropy", [logits, labels]):
            # Low confidence is given to all non-true labels, uniformly.
            low_confidence = (1.0 - confidence) / tf.to_float(vocab_size - 1)
            # Normalizing constant is the best cross-entropy value with soft targets.
            # We subtract it just for readability, makes no difference on learning.
            normalizing = -(confidence * tf.log(confidence) + tf.to_float(vocab_size - 1) * low_confidence * tf.log(low_confidence + 1e-20))
            # Soft targets.
            soft_targets = tf.one_hot(
                tf.cast(labels, tf.int32),
                depth=vocab_size,
                on_value=confidence,
                off_value=low_confidence)
            xentropy = tf.nn.softmax_cross_entropy_with_logits(
                logits=logits, labels=soft_targets)
        return xentropy - normalizing
//  [python]
def smoothing_cross_entropy(self,logits, labels, vocab_size, confidence=0.9): #confidence = 1.0 - label_smoothing. where label_smooth=0.1. from http://github.com/tensorflow/tensor2tensor
        """Cross entropy with label smoothing to limit over-confidence."""
        with tf.name_scope("smoothing_cross_entropy", [logits, labels]):
            # Low confidence is given to all non-true labels, uniformly.
            low_confidence = (1.0 - confidence) / tf.cast(vocab_size - 1, tf.float32)
            # Normalizing constant is the best cross-entropy value with soft targets.
            # We subtract it just for readability, makes no difference on learning.
            normalizing = -(confidence * tf.log(confidence) + tf.cast(vocab_size - 1, tf.float32) * low_confidence * tf.log(low_confidence + 1e-20))
            # Soft targets.
            soft_targets = tf.one_hot(
                tf.cast(labels, tf.int32),
                depth=vocab_size,
                on_value=confidence,
                off_value=low_confidence)
            xentropy = tf.nn.softmax_cross_entropy_with_logits(
                logits=logits, labels=soft_targets)
        return xentropy - normalizing
//  [python]
def learning_rate_schedule(current_epoch):
  scaled_lr = (_LEARNING_RATE *
               (FLAGS.train_batch_size * FLAGS.num_cores / 256.0))
  decay_rate = (scaled_lr * LR_SCHEDULE[0][0] *
                tf.to_float(current_epoch) / LR_SCHEDULE[0][1])
  for mult, start_epoch in LR_SCHEDULE:
    decay_rate = tf.where(current_epoch < start_epoch,
                          decay_rate, scaled_lr * mult)
  return decay_rate
//  [python]
def learning_rate_schedule(current_epoch):
  scaled_lr = (_LEARNING_RATE *
               (FLAGS.train_batch_size * FLAGS.num_cores / 256.0))
  decay_rate = (scaled_lr * LR_SCHEDULE[0][0] *
                tf.cast(current_epoch, tf.float32) / LR_SCHEDULE[0][1])
  for mult, start_epoch in LR_SCHEDULE:
    decay_rate = tf.where(current_epoch < start_epoch,
                          decay_rate, scaled_lr * mult)
  return decay_rate
//  [python]
def bytenet_internal(inputs, targets, hparams):
  """ByteNet, main step used for training."""
  with tf.variable_scope("bytenet"):
    # Flatten inputs and extend length by 50%.
    inputs = tf.expand_dims(common_layers.flatten4d3d(inputs), axis=2)
    extend_length = tf.to_int32(0.5 * tf.to_float(tf.shape(inputs)[1]))
    inputs_shape = inputs.shape.as_list()
    inputs = tf.pad(inputs, [[0, 0], [0, extend_length], [0, 0], [0, 0]])
    inputs_shape[1] = None
    inputs.set_shape(inputs_shape)  # Don't lose the other shapes when padding.
    # Pad inputs and targets to be the same length, divisible by 50.
    inputs, targets = common_layers.pad_to_same_length(
        inputs, targets, final_length_divisible_by=50)
    final_encoder = residual_dilated_conv(inputs, hparams.num_block_repeat,
                                          "SAME", "encoder", hparams)

    shifted_targets = common_layers.shift_right(targets)
    kernel = (hparams.kernel_height, hparams.kernel_width)
    decoder_start = common_layers.conv_block(
        tf.concat([final_encoder, shifted_targets], axis=3),
        hparams.hidden_size, [((1, 1), kernel)],
        padding="LEFT")

    return residual_dilated_conv(decoder_start, hparams.num_block_repeat,
                                 "LEFT", "decoder", hparams)