jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def filter_groundtruth_with_nan_box_coordinates(tensor_dict):
  """Filters out groundtruth with no bounding boxes.

  Args:
    tensor_dict: a dictionary of following groundtruth tensors -
      fields.InputDataFields.groundtruth_boxes
      fields.InputDataFields.groundtruth_classes
      fields.InputDataFields.groundtruth_keypoints
      fields.InputDataFields.groundtruth_instance_masks
      fields.InputDataFields.groundtruth_is_crowd
      fields.InputDataFields.groundtruth_area
      fields.InputDataFields.groundtruth_label_types

  Returns:
    a dictionary of tensors containing only the groundtruth that have bounding
    boxes.
  """
  groundtruth_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
  nan_indicator_vector = tf.greater(tf.reduce_sum(tf.cast(
      tf.is_nan(groundtruth_boxes), tf.int32), reduction_indices=[1]), 0)
  valid_indicator_vector = tf.logical_not(nan_indicator_vector)
  valid_indices = tf.where(valid_indicator_vector)

  return retain_groundtruth(tensor_dict, valid_indices)
//  [python]
def census_input_fn_from_csv_file(data_file, num_epochs, shuffle, batch_size):
    def _parse_csv(value):
        columns = tf.decode_csv(value, record_defaults=CENSUS_CONFIG['columns_defaults'])
        features = dict(zip(CENSUS_CONFIG['columns'],columns))
        labels = tf.equal(features.pop('income_bracket'), '>50K')
        labels = tf.reshape(labels, [-1])
        labels = tf.to_float(labels)
        return features, labels
    assert tf.io.gfile.exists(data_file), ('no file named : ' + str(data_file))
    dataset = tf.data.TextLineDataset(data_file)
    if shuffle:
        dataset = dataset.shuffle(buffer_size=5000)
    dataset = dataset.map(_parse_csv, num_parallel_calls=5)
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)
    iterator = dataset.make_one_shot_iterator()
    features, labels = iterator.get_next()
    return features, labels
//  [python]
def census_input_fn_from_csv_file(data_file, num_epochs, shuffle, batch_size):
    def _parse_csv(value):
        columns = tf.decode_csv(value, record_defaults=CENSUS_CONFIG['columns_defaults'])
        features = dict(zip(CENSUS_CONFIG['columns'],columns))
        labels = tf.equal(features.pop('income_bracket'), '>50K')
        labels = tf.reshape(labels, [-1])
        labels = tf.cast(labels, tf.float32)
        return features, labels
    assert tf.io.gfile.exists(data_file), ('no file named : ' + str(data_file))
    dataset = tf.data.TextLineDataset(data_file)
    if shuffle:
        dataset = dataset.shuffle(buffer_size=5000)
    dataset = dataset.map(_parse_csv, num_parallel_calls=5)
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)
    iterator = dataset.make_one_shot_iterator()
    features, labels = iterator.get_next()
    return features, labels
//  [python]
def q_net(x, posterior_flow, observed=None, n_z=None):
    net = spt.BayesianNet(observed=observed)

    # compute the hidden features
    with arg_scope([spt.layers.dense],
                   activation_fn=tf.nn.leaky_relu,
                   kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg)):
        h_x = tf.to_float(x)
        h_x = spt.layers.dense(h_x, 500)
        h_x = spt.layers.dense(h_x, 500)

    # sample z ~ q(z|x)
    z_mean = spt.layers.dense(h_x, config.z_dim, name='z_mean')
    z_std = 1e-4 + tf.nn.softplus(
        spt.layers.dense(h_x, config.z_dim, name='z_std'))
    z = net.add('z', spt.Normal(mean=z_mean, std=z_std), n_samples=n_z,
                group_ndims=1, flow=posterior_flow)

    return net
//  [python]
def q_net(x, posterior_flow, observed=None, n_z=None):
    net = spt.BayesianNet(observed=observed)

    # compute the hidden features
    with arg_scope([spt.layers.dense],
                   activation_fn=tf.nn.leaky_relu,
                   kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg)):
        h_x = tf.cast(x, dtype=tf.float32)
        h_x = spt.layers.dense(h_x, 500)
        h_x = spt.layers.dense(h_x, 500)

    # sample z ~ q(z|x)
    z_mean = spt.layers.dense(h_x, config.z_dim, name='z_mean')
    z_std = 1e-4 + tf.nn.softplus(
        spt.layers.dense(h_x, config.z_dim, name='z_std'))
    z = net.add('z', spt.Normal(mean=z_mean, std=z_std), n_samples=n_z,
                group_ndims=1, flow=posterior_flow)

    return net
//  [python]
def random_flip_left_right(image, masks, boxes, keypoints):

    def flip(image, masks, boxes, keypoints):

        flipped_image = tf.image.flip_left_right(image)
        flipped_masks = tf.image.flip_left_right(masks)

        y, x, v = tf.unstack(keypoints, axis=2)
        width = tf.shape(image)[1]
        flipped_x = width - 1 - x
        flipped_keypoints = tf.stack([y, flipped_x, v], axis=2)

        width = tf.to_float(width)
        ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=1)
        flipped_boxes = tf.stack([ymin, width - xmax, ymax, width - xmin], axis=1)

        """
        The keypoint order:
        0: 'nose',
        1: 'left eye', 2: 'right eye',
        3: 'left ear', 4: 'right ear',
        5: 'left shoulder', 6: 'right shoulder',
        7: 'left elbow', 8: 'right elbow',
        9: 'left wrist', 10: 'right wrist',
        11: 'left hip', 12: 'right hip',
        13: 'left knee', 14: 'right knee',
        15: 'left ankle', 16: 'right ankle'
        """

        correct_order = tf.constant([0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15])
        flipped_keypoints = tf.gather(flipped_keypoints, correct_order, axis=1)

        return flipped_image, flipped_masks, flipped_boxes, flipped_keypoints

    do_it = tf.less(tf.random_uniform([]), 0.5)
    image, masks, boxes, keypoints = tf.cond(
        do_it,
        lambda: flip(image, masks, boxes, keypoints),
        lambda: (image, masks, boxes, keypoints)
    )
    return image, masks, boxes, keypoints
//  [python]
def random_flip_left_right(image, masks, boxes, keypoints):

    def flip(image, masks, boxes, keypoints):

        flipped_image = tf.image.flip_left_right(image)
        flipped_masks = tf.image.flip_left_right(masks)

        y, x, v = tf.unstack(keypoints, axis=2)
        width = tf.shape(image)[1]
        flipped_x = width - 1 - x
        flipped_keypoints = tf.stack([y, flipped_x, v], axis=2)

        width = tf.cast(width, tf.float32)
        ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=1)
        flipped_boxes = tf.stack([ymin, width - xmax, ymax, width - xmin], axis=1)

        """
        The keypoint order:
        0: 'nose',
        1: 'left eye', 2: 'right eye',
        3: 'left ear', 4: 'right ear',
        5: 'left shoulder', 6: 'right shoulder',
        7: 'left elbow', 8: 'right elbow',
        9: 'left wrist', 10: 'right wrist',
        11: 'left hip', 12: 'right hip',
        13: 'left knee', 14: 'right knee',
        15: 'left ankle', 16: 'right ankle'
        """

        correct_order = tf.constant([0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15])
        flipped_keypoints = tf.gather(flipped_keypoints, correct_order, axis=1)

        return flipped_image, flipped_masks, flipped_boxes, flipped_keypoints

    do_it = tf.less(tf.random_uniform([]), 0.5)
    image, masks, boxes, keypoints = tf.cond(
        do_it,
        lambda: flip(image, masks, boxes, keypoints),
        lambda: (image, masks, boxes, keypoints)
    )
    return image, masks, boxes, keypoints
//  [python]
def parse(self, example_proto):
        """
        Returns:
            image: a float tensor with shape [height, width, 3],
                an RGB image with pixel values in the range [0, 1].
            masks: a float tensor with shape [height / DOWNSAMPLE, width / DOWNSAMPLE, 2].
            boxes: a float tensor with shape [num_persons, 4], in absolute coordinates.
            keypoints: an int tensor with shape [num_persons, 17, 3], in absolute coordinates.
        """
        features = {
            'image': tf.FixedLenFeature([], tf.string),
            'num_persons': tf.FixedLenFeature([], tf.int64),
            'boxes': tf.FixedLenSequenceFeature([], tf.float32, allow_missing=True),
            'keypoints': tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
            'masks': tf.FixedLenFeature([], tf.string)
        }
        parsed_features = tf.parse_single_example(example_proto, features)

        # get an image
        image = tf.image.decode_jpeg(parsed_features['image'], channels=3)
        image = tf.image.convert_image_dtype(image, tf.float32)
        # now pixel values are scaled to the [0, 1] range

        # get number of people on the image
        num_persons = tf.to_int32(parsed_features['num_persons'])
        # it is assumed that num_persons > 0

        # get groundtruth boxes, they are in absolute coordinates
        boxes = tf.reshape(parsed_features['boxes'], [num_persons, 4])
        # they are used to guide the data augmentation (when doing a random crop)
        # and to choose sigmas for gaussian blobs

        # get keypoints, they are in absolute coordinates
        keypoints = tf.to_int32(parsed_features['keypoints'])
        keypoints = tf.reshape(keypoints, [num_persons, 17, 3])

        # get size of masks, they are downsampled
        shape = tf.shape(image)
        image_height, image_width = shape[0], shape[1]
        masks_height = tf.to_int32(tf.ceil(image_height/DOWNSAMPLE))
        masks_width = tf.to_int32(tf.ceil(image_width/DOWNSAMPLE))
        # (we use the 'SAME' padding in the networks)

        # get masks (loss and segmentation masks)
        masks = tf.decode_raw(parsed_features['masks'], tf.uint8)
        # unpack bits (reverse np.packbits)
        b = tf.constant([128, 64, 32, 16, 8, 4, 2, 1], dtype=tf.uint8)
        masks = tf.reshape(tf.bitwise.bitwise_and(masks[:, None], b), [-1])
        masks = masks[:(masks_height * masks_width * 2)]
        masks = tf.cast(masks > 0, tf.uint8)

        # reshape to the initial form
        masks = tf.reshape(masks, [masks_height, masks_width, 2])
        masks = tf.to_float(masks)  # it has binary values only

        return image, masks, boxes, keypoints
//  [python]
def parse(self, example_proto):
        """
        Returns:
            image: a float tensor with shape [height, width, 3],
                an RGB image with pixel values in the range [0, 1].
            masks: a float tensor with shape [height / DOWNSAMPLE, width / DOWNSAMPLE, 2].
            boxes: a float tensor with shape [num_persons, 4], in absolute coordinates.
            keypoints: an int tensor with shape [num_persons, 17, 3], in absolute coordinates.
        """
        features = {
            'image': tf.FixedLenFeature([], tf.string),
            'num_persons': tf.FixedLenFeature([], tf.int64),
            'boxes': tf.FixedLenSequenceFeature([], tf.float32, allow_missing=True),
            'keypoints': tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
            'masks': tf.FixedLenFeature([], tf.string)
        }
        parsed_features = tf.parse_single_example(example_proto, features)

        # get an image
        image = tf.image.decode_jpeg(parsed_features['image'], channels=3)
        image = tf.image.convert_image_dtype(image, tf.float32)
        # now pixel values are scaled to the [0, 1] range

        # get number of people on the image
        num_persons = tf.cast(parsed_features['num_persons'], tf.int32)
        # it is assumed that num_persons > 0

        # get groundtruth boxes, they are in absolute coordinates
        boxes = tf.reshape(parsed_features['boxes'], [num_persons, 4])
        # they are used to guide the data augmentation (when doing a random crop)
        # and to choose sigmas for gaussian blobs

        # get keypoints, they are in absolute coordinates
        keypoints = tf.cast(parsed_features['keypoints'], tf.int32)
        keypoints = tf.reshape(keypoints, [num_persons, 17, 3])

        # get size of masks, they are downsampled
        shape = tf.shape(image)
        image_height, image_width = shape[0], shape[1]
        masks_height = tf.cast(tf.ceil(image_height/DOWNSAMPLE), tf.int32)
        masks_width = tf.cast(tf.ceil(image_width/DOWNSAMPLE), tf.int32)
        # (we use the 'SAME' padding in the networks)

        # get masks (loss and segmentation masks)
        masks = tf.decode_raw(parsed_features['masks'], tf.uint8)
        # unpack bits (reverse np.packbits)
        b = tf.constant([128, 64, 32, 16, 8, 4, 2, 1], dtype=tf.uint8)
        masks = tf.reshape(tf.bitwise.bitwise_and(masks[:, None], b), [-1])
        masks = masks[:(masks_height * masks_width * 2)]
        masks = tf.cast(masks > 0, tf.uint8)

        # reshape to the initial form
        masks = tf.reshape(masks, [masks_height, masks_width, 2])
        masks = tf.cast(masks, tf.float32)  # it has binary values only

        return image, masks, boxes, keypoints
//  [python]
def census_input_fn_from_tfrecords(data_file, num_epochs, shuffle, batch_size):
    def _parse_census_TFRecords_fn(record):
        features = {
            # int
            'age':            tf.io.FixedLenFeature([], tf.int64),
            'fnlwgt':         tf.io.FixedLenFeature([], tf.int64),
            'education_num':  tf.io.FixedLenFeature([], tf.int64),
            'capital_gain':   tf.io.FixedLenFeature([], tf.int64),
            'capital_loss':   tf.io.FixedLenFeature([], tf.int64),
            'hours_per_week': tf.io.FixedLenFeature([], tf.int64),
            # string
            'gender':         tf.io.FixedLenFeature([], tf.string),
            'education':      tf.io.FixedLenFeature([], tf.string),
            'marital_status': tf.io.FixedLenFeature([], tf.string),
            'relationship':   tf.io.FixedLenFeature([], tf.string),
            'race':           tf.io.FixedLenFeature([], tf.string),
            'workclass':      tf.io.FixedLenFeature([], tf.string),
            'native_country': tf.io.FixedLenFeature([], tf.string),
            'occupation':     tf.io.FixedLenFeature([], tf.string),
            'income_bracket': tf.io.FixedLenFeature([], tf.string),
        }
        features = tf.io.parse_single_example(record, features)
        labels = tf.equal(features.pop('income_bracket'), '>50K')
        labels = tf.reshape(labels, [-1])
        labels = tf.to_float(labels)
        return features, labels
    assert tf.io.gfile.exists(data_file), ('no file named: ' + str(data_file))
    dataset = tf.data.TFRecordDataset(data_file).map(_parse_census_TFRecords_fn, num_parallel_calls=10)
    if shuffle:
        dataset = dataset.shuffle(buffer_size=5000)
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)
    iterator = dataset.make_one_shot_iterator()
    features, labels = iterator.get_next()
    return features, labels
//  [python]
def census_input_fn_from_tfrecords(data_file, num_epochs, shuffle, batch_size):
    def _parse_census_TFRecords_fn(record):
        features = {
            # int
            'age':            tf.io.FixedLenFeature([], tf.int64),
            'fnlwgt':         tf.io.FixedLenFeature([], tf.int64),
            'education_num':  tf.io.FixedLenFeature([], tf.int64),
            'capital_gain':   tf.io.FixedLenFeature([], tf.int64),
            'capital_loss':   tf.io.FixedLenFeature([], tf.int64),
            'hours_per_week': tf.io.FixedLenFeature([], tf.int64),
            # string
            'gender':         tf.io.FixedLenFeature([], tf.string),
            'education':      tf.io.FixedLenFeature([], tf.string),
            'marital_status': tf.io.FixedLenFeature([], tf.string),
            'relationship':   tf.io.FixedLenFeature([], tf.string),
            'race':           tf.io.FixedLenFeature([], tf.string),
            'workclass':      tf.io.FixedLenFeature([], tf.string),
            'native_country': tf.io.FixedLenFeature([], tf.string),
            'occupation':     tf.io.FixedLenFeature([], tf.string),
            'income_bracket': tf.io.FixedLenFeature([], tf.string),
        }
        features = tf.io.parse_single_example(record, features)
        labels = tf.equal(features.pop('income_bracket'), '>50K')
        labels = tf.reshape(labels, [-1])
        labels = tf.cast(labels, tf.float32)
        return features, labels
    assert tf.io.gfile.exists(data_file), ('no file named: ' + str(data_file))
    dataset = tf.data.TFRecordDataset(data_file).map(_parse_census_TFRecords_fn, num_parallel_calls=10)
    if shuffle:
        dataset = dataset.shuffle(buffer_size=5000)
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)
    iterator = dataset.make_one_shot_iterator()
    features, labels = iterator.get_next()
    return features, labels
//  [python]
def softmax_summaries(loss, logits, one_hot_labels, name="softmax"):
  """Create the softmax summaries for this cross entropy loss.

  Args:
    loss: Cross-entropy loss.
    logits: The [batch_size, classes] float tensor representing the logits.
    one_hot_labels: The float tensor representing actual class ids. If this is
      [batch_size, classes], then we take the argmax of it first.
    name: Prepended to summary scope.
  """
  tf.summary.scalar(name + "_loss", loss)

  one_hot_labels = tf.cond(
      tf.equal(tf.rank(one_hot_labels),
               2), lambda: tf.to_int32(tf.argmax(one_hot_labels, 1)),
      lambda: tf.to_int32(one_hot_labels))

  in_top_1 = tf.nn.in_top_k(logits, one_hot_labels, 1)
  tf.summary.scalar(name + "_precision@1",
                    tf.reduce_mean(tf.to_float(in_top_1)))
  in_top_5 = tf.nn.in_top_k(logits, one_hot_labels, 5)
  tf.summary.scalar(name + "_precision@5",
                    tf.reduce_mean(tf.to_float(in_top_5)))
//  [python]
def softmax_summaries(loss, logits, one_hot_labels, name="softmax"):
  """Create the softmax summaries for this cross entropy loss.

  Args:
    loss: Cross-entropy loss.
    logits: The [batch_size, classes] float tensor representing the logits.
    one_hot_labels: The float tensor representing actual class ids. If this is
      [batch_size, classes], then we take the argmax of it first.
    name: Prepended to summary scope.
  """
  tf.summary.scalar(name + "_loss", loss)

  one_hot_labels = tf.cond(
      tf.equal(tf.rank(one_hot_labels),
               2), lambda: tf.cast(tf.argmax(one_hot_labels, 1), tf.int32),
      lambda: tf.cast(one_hot_labels, tf.int32))

  in_top_1 = tf.nn.in_top_k(logits, one_hot_labels, 1)
  tf.summary.scalar(name + "_precision@1",
                    tf.reduce_mean(tf.to_float(in_top_1)))
  in_top_5 = tf.nn.in_top_k(logits, one_hot_labels, 5)
  tf.summary.scalar(name + "_precision@5",
                    tf.reduce_mean(tf.to_float(in_top_5)))
//  [python]
def parse_example(serialized_example):
  """Parse example."""
  features = tf.parse_single_example(
      serialized_example,
      features={
          "question": tf.FixedLenFeature([], tf.string),
          "context": tf.FixedLenFeature([], tf.string),
          "answer_start": tf.FixedLenFeature([], tf.int64),
          "answer_end": tf.FixedLenFeature([], tf.int64),
      })
  features["question"] = split_on_whitespace(features["question"])
  features["context"] = split_on_whitespace(features["context"])
  features["answer_start"] = tf.to_int32(features["answer_start"])
  features["answer_end"] = tf.to_int32(features["answer_end"])
  return features
//  [python]
def parse_example(serialized_example):
  """Parse example."""
  features = tf.parse_single_example(
      serialized_example,
      features={
          "question": tf.FixedLenFeature([], tf.string),
          "context": tf.FixedLenFeature([], tf.string),
          "answer_start": tf.FixedLenFeature([], tf.int64),
          "answer_end": tf.FixedLenFeature([], tf.int64),
      })
  features["question"] = split_on_whitespace(features["question"])
  features["context"] = split_on_whitespace(features["context"])
  features["answer_start"] = tf.cast(features["answer_start"], tf.int32)
  features["answer_end"] = tf.cast(features["answer_end"], tf.int32)
  return features
//  [python]
def compute_null_weights(labels, null_weight):
  has_real_context = tf.to_float(tf.reduce_any(tf.greater(labels, 0), 1))
  weights = has_real_context + (1.0 - has_real_context) * null_weight
  return weights
//  [python]
def compute_null_weights(labels, null_weight):
  has_real_context = tf.cast(tf.reduce_any(tf.greater(labels, 0), 1), tf.float32)
  weights = has_real_context + (1.0 - has_real_context) * null_weight
  return weights
//  [python]
def compute_thresholded_labels(labels, null_threshold=4):
  """Computes thresholded labels.

  Args:
    labels: <int32> [batch_size, num_annotators]
    null_threshold: If number of null annotations is greater than or equal to
      this threshold, all annotations are set to null for this example.

  Returns:
    thresholded_labels: <int32> [batch_size, num_annotators]
  """
  null_labels = tf.equal(labels, 0)

  # <int32> [batch_size]
  null_count = tf.reduce_sum(tf.to_int32(null_labels), 1)
  threshold_mask = tf.less(null_count, null_threshold)

  # <bool> [batch_size, num_annotators]
  threshold_mask = tf.tile(
      tf.expand_dims(threshold_mask, -1), [1, tf.shape(labels)[1]])

  # <bool> [batch_size, num_annotators]
  thresholded_labels = tf.where(
      threshold_mask, x=labels, y=tf.zeros_like(labels))
  return thresholded_labels
//  [python]
def compute_thresholded_labels(labels, null_threshold=4):
  """Computes thresholded labels.

  Args:
    labels: <int32> [batch_size, num_annotators]
    null_threshold: If number of null annotations is greater than or equal to
      this threshold, all annotations are set to null for this example.

  Returns:
    thresholded_labels: <int32> [batch_size, num_annotators]
  """
  null_labels = tf.equal(labels, 0)

  # <int32> [batch_size]
  null_count = tf.reduce_sum(tf.cast(null_labels, tf.int32), 1)
  threshold_mask = tf.less(null_count, null_threshold)

  # <bool> [batch_size, num_annotators]
  threshold_mask = tf.tile(
      tf.expand_dims(threshold_mask, -1), [1, tf.shape(labels)[1]])

  # <bool> [batch_size, num_annotators]
  thresholded_labels = tf.where(
      threshold_mask, x=labels, y=tf.zeros_like(labels))
  return thresholded_labels
//  [python]
def truncate_contexts(context_token_ids, num_contexts, context_len,
                      max_contexts, max_context_len):
  """Truncate context based on max_contexts and max_context_words.

  Drops more than `max_contexts` contexts entirely, and truncates
  each remaining context to contain at most `max_context_words` words.

  Args:
    context_token_ids: <int32> [batch_size, padded_num_contexts,
      padded_context_len].
    num_contexts: <int32> [batch_size] for number of contexts in each batch.
    context_len: <int32> [batch_size, padded_num_contexts] for length of each
      context.
    max_contexts: <int32> for maximum number of contexts per example.
    max_context_len: <int32> for maximum length of each context.

  Returns:
    Truncated tensors for context_token_ids, context_len and document_len
    based on max_contexts and max_context_words.
  """
  # Compute effective max contexts and max context words.
  padded_num_contexts = tf.shape(context_token_ids)[1]
  padded_context_len = tf.shape(context_token_ids)[2]

  max_contexts = tf.to_int32(tf.minimum(padded_num_contexts, max_contexts))
  max_context_len = tf.to_int32(tf.minimum(padded_context_len, max_context_len))

  # Truncate context_token_ids, context_len, and document len
  context_token_ids = context_token_ids[:, :max_contexts, :max_context_len]
  context_len = context_len[:, :max_contexts]
  context_len = tf.minimum(context_len, max_context_len)
  num_contexts = tf.minimum(num_contexts, max_contexts)

  return context_token_ids, num_contexts, context_len
//  [python]
def truncate_contexts(context_token_ids, num_contexts, context_len,
                      max_contexts, max_context_len):
  """Truncate context based on max_contexts and max_context_words.

  Drops more than `max_contexts` contexts entirely, and truncates
  each remaining context to contain at most `max_context_words` words.

  Args:
    context_token_ids: <int32> [batch_size, padded_num_contexts,
      padded_context_len].
    num_contexts: <int32> [batch_size] for number of contexts in each batch.
    context_len: <int32> [batch_size, padded_num_contexts] for length of each
      context.
    max_contexts: <int32> for maximum number of contexts per example.
    max_context_len: <int32> for maximum length of each context.

  Returns:
    Truncated tensors for context_token_ids, context_len and document_len
    based on max_contexts and max_context_words.
  """
  # Compute effective max contexts and max context words.
  padded_num_contexts = tf.shape(context_token_ids)[1]
  padded_context_len = tf.shape(context_token_ids)[2]

  max_contexts = tf.cast(tf.minimum(padded_num_contexts, max_contexts), tf.int32)
  max_context_len = tf.cast(tf.minimum(padded_context_len, max_context_len), tf.int32)

  # Truncate context_token_ids, context_len, and document len
  context_token_ids = context_token_ids[:, :max_contexts, :max_context_len]
  context_len = context_len[:, :max_contexts]
  context_len = tf.minimum(context_len, max_context_len)
  num_contexts = tf.minimum(num_contexts, max_contexts)

  return context_token_ids, num_contexts, context_len
//  [python]
def compute_match_stats(predictions, labels):
  """Compute statistics that are used to compute evaluation metrics.

  Args:
    predictions: <int32> [batch_size]
    labels: <int32> [batch_size, num_annotators]

  Returns:
    numerator: <float> [batch_size]
    recall_weights: <float> [batch_size]
    precision_weights: <float> [batch_size]
  """
  # <int32> [batch_size, num_labels]
  thresholded_labels = compute_thresholded_labels(labels)
  non_null_mask = tf.greater(thresholded_labels, 0)

  # <bool> [batch_size, num_labels]
  exact_match = tf.equal(tf.expand_dims(predictions, -1), thresholded_labels)
  non_null_match = tf.logical_and(exact_match, non_null_mask)

  # <float> [batch_size]
  non_null_match = tf.to_float(tf.reduce_any(non_null_match, axis=1))
  non_null_gold = tf.to_float(tf.reduce_any(non_null_mask, 1))
  non_null_predictions = tf.to_float(tf.greater(predictions, 0))

  return non_null_match, non_null_gold, non_null_predictions
//  [python]
def compute_match_stats(predictions, labels):
  """Compute statistics that are used to compute evaluation metrics.

  Args:
    predictions: <int32> [batch_size]
    labels: <int32> [batch_size, num_annotators]

  Returns:
    numerator: <float> [batch_size]
    recall_weights: <float> [batch_size]
    precision_weights: <float> [batch_size]
  """
  # <int32> [batch_size, num_labels]
  thresholded_labels = compute_thresholded_labels(labels)
  non_null_mask = tf.greater(thresholded_labels, 0)

  # <bool> [batch_size, num_labels]
  exact_match = tf.equal(tf.expand_dims(predictions, -1), thresholded_labels)
  non_null_match = tf.logical_and(exact_match, non_null_mask)

  # <float> [batch_size]
  non_null_match = tf.cast(tf.reduce_any(non_null_match, axis=1), tf.float32)
  non_null_gold = tf.cast(tf.reduce_any(non_null_mask, 1), tf.float32)
  non_null_predictions = tf.cast(tf.greater(predictions, 0), tf.float32)

  return non_null_match, non_null_gold, non_null_predictions
//  [python]
def _seconds_to_internal_time(seconds):
  """Converts seconds to fixed-precision time."""
  return tf.to_int64(tf.round(seconds * _INTERNAL_TIME_PRECISION))
//  [python]
def _seconds_to_internal_time(seconds):
  """Converts seconds to fixed-precision time."""
  return tf.cast(tf.round(seconds * _INTERNAL_TIME_PRECISION), tf.int64)
//  [python]
def _internal_time_to_seconds(internal_time):
  """Converts fixed-precision time to seconds."""
  return tf.to_float(internal_time / _INTERNAL_TIME_PRECISION)
//  [python]
def _internal_time_to_seconds(internal_time):
  """Converts fixed-precision time to seconds."""
  return tf.cast(internal_time / _INTERNAL_TIME_PRECISION, tf.float32)
//  [python]
def _compute_loss(self, logits):
    """Compute optimization loss."""
    target_output = self.iterator.target_output
    if self.time_major:
      target_output = tf.transpose(target_output)
    max_time = self.get_max_time(target_output)
    crossent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_output, logits=logits)
    target_weights = tf.sequence_mask(
        self.iterator.target_sequence_length, max_time, dtype=logits.dtype)
    if self.time_major:
      target_weights = tf.transpose(target_weights)

    loss = tf.reduce_sum(
        crossent * target_weights) / tf.to_float(self.batch_size)
    return loss
//  [python]
def _compute_loss(self, logits):
    """Compute optimization loss."""
    target_output = self.iterator.target_output
    if self.time_major:
      target_output = tf.transpose(target_output)
    max_time = self.get_max_time(target_output)
    crossent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_output, logits=logits)
    target_weights = tf.sequence_mask(
        self.iterator.target_sequence_length, max_time, dtype=logits.dtype)
    if self.time_major:
      target_weights = tf.transpose(target_weights)

    loss = tf.reduce_sum(
        crossent * target_weights) / tf.cast(self.batch_size, dtype=tf.float32)
    return loss
//  [python]
def input_fn_builder(
    input_file, seq_length, is_training, drop_remainder, use_tpu, bsz, is_v2
):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    name_to_features = {
        'unique_ids': tf.FixedLenFeature([], tf.int64),
        'input_ids': tf.FixedLenFeature([seq_length], tf.int64),
        'input_mask': tf.FixedLenFeature([seq_length], tf.int64),
        'segment_ids': tf.FixedLenFeature([seq_length], tf.int64),
    }
    # p_mask is not required for SQuAD v1.1
    if is_v2:
        name_to_features['p_mask'] = tf.FixedLenFeature([seq_length], tf.int64)

    if is_training:
        name_to_features['start_positions'] = tf.FixedLenFeature([], tf.int64)
        name_to_features['end_positions'] = tf.FixedLenFeature([], tf.int64)
        name_to_features['is_impossible'] = 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."""
        if use_tpu:
            batch_size = params['batch_size']
        else:
            batch_size = bsz

        # 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(
            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 input_fn_builder(
    input_file, seq_length, is_training, drop_remainder, use_tpu, bsz, is_v2
):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    name_to_features = {
        'unique_ids': tf.FixedLenFeature([], tf.int64),
        'input_ids': tf.FixedLenFeature([seq_length], tf.int64),
        'input_mask': tf.FixedLenFeature([seq_length], tf.int64),
        'segment_ids': tf.FixedLenFeature([seq_length], tf.int64),
    }
    # p_mask is not required for SQuAD v1.1
    if is_v2:
        name_to_features['p_mask'] = tf.FixedLenFeature([seq_length], tf.int64)

    if is_training:
        name_to_features['start_positions'] = tf.FixedLenFeature([], tf.int64)
        name_to_features['end_positions'] = tf.FixedLenFeature([], tf.int64)
        name_to_features['is_impossible'] = 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."""
        if use_tpu:
            batch_size = params['batch_size']
        else:
            batch_size = bsz

        # 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(
            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 testGreedyWithCornerCase(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 3
    decode_length = 2

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[0.2, 0.1, 0.7], [0.4, 0.1, 0.5]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    self.assertAllEqual([[[0, 2, 2]]], ids)
    self.assertAllClose([[0.7 * 0.5]], np.exp(probs))
//  [python]
def testGreedyWithCornerCase(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 3
    decode_length = 2

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[0.2, 0.1, 0.7], [0.4, 0.1, 0.5]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.cast(tf.log(probabilities[pos - 1, :]), tf.float32)
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    self.assertAllEqual([[[0, 2, 2]]], ids)
    self.assertAllClose([[0.7 * 0.5]], np.exp(probs))
//  [python]
def testNotGreedyBeamTwoWithoutStopEarly(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        stop_early=False)

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    # given stop_early = False, the algorithm will return all the beams
    # so we can test all of them here
    self.assertAllEqual([[[0, 2, 1, 0], [0, 2, 0, 1]]], ids)
    self.assertAllClose([[0.8 * 0.5, 0.8 * 0.4 * 0.9]], np.exp(probs))
//  [python]
def testNotGreedyBeamTwoWithoutStopEarly(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.cast(tf.log(probabilities[pos - 1, :]), dtype=tf.float32)
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        stop_early=False)

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    # given stop_early = False, the algorithm will return all the beams
    # so we can test all of them here
    self.assertAllEqual([[[0, 2, 1, 0], [0, 2, 0, 1]]], ids)
    self.assertAllClose([[0.8 * 0.5, 0.8 * 0.4 * 0.9]], np.exp(probs))
//  [python]
def testNotGreedyBeamTwoWithStopEarly(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        stop_early=True)  # default value, but just to make this explicit

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    # given stop_early = True, the only 'assurance' is w.r.t. the first beam
    # (i.e., other beams may not even be completed)
    # so, we check only the first beam
    first_beam = ids[:, 0]
    first_probs = probs[:, 0]
    self.assertAllEqual([[0, 2, 1]], first_beam)
    self.assertAllClose([0.8 * 0.5], np.exp(first_probs))
//  [python]
def testNotGreedyBeamTwoWithStopEarly(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.cast(tf.log(probabilities[pos - 1, :]), tf.float32)
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        stop_early=True)  # default value, but just to make this explicit

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    # given stop_early = True, the only 'assurance' is w.r.t. the first beam
    # (i.e., other beams may not even be completed)
    # so, we check only the first beam
    first_beam = ids[:, 0]
    first_probs = probs[:, 0]
    self.assertAllEqual([[0, 2, 1]], first_beam)
    self.assertAllClose([0.8 * 0.5], np.exp(first_probs))
//  [python]
def testGreedyBatchOne(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 2
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO

    # Test that beam search finds the most probable sequence.
    # These probabilities represent the following search
    #
    #               G0 (0)
    #                  / \
    #                /     \
    #              /         \
    #            /             \
    #         0(0.7)          1(0.3)
    #           / \
    #          /   \
    #         /     \
    #     0(0.4) 1(0.6)
    #        /\
    #       /  \
    #      /    \
    #    0(0.5) 1(0.5)
    # and the following decoding probabilities
    # 0000 - 0.7 * 0.4  * 0.1
    # 0001 - 0.7 * 0.4  * 0.9
    # 001 - 0.7 * 0.6 (Best)
    # 01 = 0.3
    #
    # 001 is the most likely sequence under these probabilities.
    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    self.assertAllEqual([[[0, 0, 1]]], ids)
    self.assertAllClose([[0.7 * 0.6]], np.exp(probs))
//  [python]
def testGreedyBatchOne(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 2
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO

    # Test that beam search finds the most probable sequence.
    # These probabilities represent the following search
    #
    #               G0 (0)
    #                  / \
    #                /     \
    #              /         \
    #            /             \
    #         0(0.7)          1(0.3)
    #           / \
    #          /   \
    #         /     \
    #     0(0.4) 1(0.6)
    #        /\
    #       /  \
    #      /    \
    #    0(0.5) 1(0.5)
    # and the following decoding probabilities
    # 0000 - 0.7 * 0.4  * 0.1
    # 0001 - 0.7 * 0.4  * 0.9
    # 001 - 0.7 * 0.6 (Best)
    # 01 = 0.3
    #
    # 001 is the most likely sequence under these probabilities.
    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.cast(tf.log(probabilities[pos - 1, :]), tf.float32)
      return logits

    final_ids, final_probs, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      probs = final_probs.eval()
    self.assertAllEqual([[[0, 0, 1]]], ids)
    self.assertAllClose([[0.7 * 0.6]], np.exp(probs))
//  [python]
def body(self, features, decode_step=None, cache=None, decoding_stats=None):
    targets = tf.to_int32(features["targets"])
    if self.is_decode:
      targets = self.process_partial_targets_decoding(targets)
    decoder_input = self.prepare_decoder(targets)
    extra_losses = []

    if not self.hparams.unconditional:  # condition on class label
      if not self.is_inputs_class_label:
        raise ValueError("SparseImagetransformer can only condition on "
                         "'inputs' feature if it represents class label.")
      inputs = features["inputs"]

      # Embed class here rather than in bottom().
      if inputs.dtype not in [tf.int32, tf.int64]:
        raise ValueError("Do not embed 'inputs' before body(). "
                         "Found dtype=%s." % inputs.dtype)
      inputs = utils.get_embeddings(
          targets=inputs,
          vocab_size=self.inputs_vocab_size,
          hidden_size=self.hidden_size,
          name="class_conditional_embedding")

      # Add class embedding to each spatial location.
      batch_size = tf.shape(targets)[0]
      hidden_size = tf.shape(inputs)[-1]
      num_middle_dims = len(decoder_input.shape) - 2
      decoder_input += tf.reshape(inputs, [batch_size] + [1] * num_middle_dims +
                                  [hidden_size])

    decoder_output = utils.transformer_decoder_layers(
        inputs=decoder_input,
        num_layers=self.num_decoder_layers,
        hparams=self.hparams,
        decode_step=decode_step,
        losses=extra_losses,
        cache=cache,
        name="decoder",
        decoding_stats=decoding_stats)
    logits = self.produce_output(decoder_output)

    # Return logits as-is in decoding mode
    if self.is_decode:
      return logits

    # Produce a summary of the output.
    results = self.multinomial_squeeze(logits, self.hparams.sampling_temp)
    results = tf.reshape(
        results, [-1, self.frame_height, self.frame_width, self.num_channels])
    if utils.is_xla_compiled():
      _IMGS["predictions"] = results

    # Prepare loss.
    loss_dict = {}
    if extra_losses:
      loss_dict["extra_loss"] = tf.add_n(extra_losses)
    return logits, loss_dict
//  [python]
def body(self, features, decode_step=None, cache=None, decoding_stats=None):
    targets = tf.cast(features["targets"], tf.int32)
    if self.is_decode:
      targets = self.process_partial_targets_decoding(targets)
    decoder_input = self.prepare_decoder(targets)
    extra_losses = []

    if not self.hparams.unconditional:  # condition on class label
      if not self.is_inputs_class_label:
        raise ValueError("SparseImagetransformer can only condition on "
                         "'inputs' feature if it represents class label.")
      inputs = features["inputs"]

      # Embed class here rather than in bottom().
      if inputs.dtype not in [tf.int32, tf.int64]:
        raise ValueError("Do not embed 'inputs' before body(). "
                         "Found dtype=%s." % inputs.dtype)
      inputs = utils.get_embeddings(
          targets=inputs,
          vocab_size=self.inputs_vocab_size,
          hidden_size=self.hidden_size,
          name="class_conditional_embedding")

      # Add class embedding to each spatial location.
      batch_size = tf.shape(targets)[0]
      hidden_size = tf.shape(inputs)[-1]
      num_middle_dims = len(decoder_input.shape) - 2
      decoder_input += tf.reshape(inputs, [batch_size] + [1] * num_middle_dims +
                                  [hidden_size])

    decoder_output = utils.transformer_decoder_layers(
        inputs=decoder_input,
        num_layers=self.num_decoder_layers,
        hparams=self.hparams,
        decode_step=decode_step,
        losses=extra_losses,
        cache=cache,
        name="decoder",
        decoding_stats=decoding_stats)
    logits = self.produce_output(decoder_output)

    # Return logits as-is in decoding mode
    if self.is_decode:
      return logits

    # Produce a summary of the output.
    results = self.multinomial_squeeze(logits, self.hparams.sampling_temp)
    results = tf.reshape(
        results, [-1, self.frame_height, self.frame_width, self.num_channels])
    if utils.is_xla_compiled():
      _IMGS["predictions"] = results

    # Prepare loss.
    loss_dict = {}
    if extra_losses:
      loss_dict["extra_loss"] = tf.add_n(extra_losses)
    return logits, loss_dict
//  [python]
def testNotGreedyBeamTwoWithAlpha(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    # Probabilities for position * batch * beam * vocab
    # Probabilities have been set such that with alpha = 3.5, the less probable
    # but longer sequence will have a better score that the shorter sequence
    # with higher log prob.
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
      return logits

    # Disable early stopping
    final_ids, final_scores, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        3.5,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      scores = final_scores.eval()
    self.assertAllClose([[
        np.log(0.8 * 0.4 * 0.9) / (8. / 6.)**3.5,
        np.log(0.8 * 0.5) / (7. / 6.)**3.5
    ]], scores)
    self.assertAllEqual([[[0, 2, 0, 1], [0, 2, 1, 0]]], ids)