jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def rotate_image_tensor(image, angle):
    s = tf.shape(image)
    image_center = tf.to_float(s[:2]) / 2.

    # Coordinates of new image
    xs, ys = tf.meshgrid(tf.range(0.,tf.to_float(s[1])), tf.range(0., tf.to_float(s[0])))
    coords_new = tf.reshape(tf.stack([ys,xs], 2), [-1, 2])

    # center coordinates since rotation center is supposed to be in the image center
    coords_new_centered = tf.to_float(coords_new) - image_center

    # Perform backward transformation of the image coordinates
    rot_mat_inv = tf.stack(
        [tf.cos(angle), tf.sin(angle), -tf.sin(angle), tf.cos(angle)])
    rot_mat_inv = tf.reshape(rot_mat_inv, shape=[2, 2])
    coord_old_centered = tf.matmul(
        rot_mat_inv, tf.transpose(coords_new_centered))
    coord_old = tf.cast(tf.round(
        tf.transpose(coord_old_centered) + image_center), tf.int32)


    # Find nearest neighbor in old image
    coord_old_y, coord_old_x = tf.unstack(coord_old, axis=1)


    # Clip values to stay inside image coordinates
    outside_y = tf.logical_or(tf.greater(
        coord_old_y, s[0]-1), tf.less(coord_old_y, 0))
    outside_x = tf.logical_or(tf.greater(
        coord_old_x, s[1]-1), tf.less(coord_old_x, 0))
    outside_ind = tf.logical_or(outside_y, outside_x)


    inside_mask = tf.logical_not(outside_ind)
    inside_mask = tf.tile(tf.reshape(inside_mask, s[:2])[...,None], tf.stack([1,1,s[2]]))

    coord_old_y = tf.clip_by_value(coord_old_y, 0, s[0]-1)
    coord_old_x = tf.clip_by_value(coord_old_x, 0, s[1]-1)
    coord_flat = coord_old_y * s[1] + coord_old_x

    im_flat = tf.reshape(image, tf.stack([-1, s[2]]))
    rot_image = tf.gather(im_flat, coord_flat)
    rot_image = tf.reshape(rot_image, s)


    return tf.where(inside_mask, rot_image, tf.zeros_like(rot_image))
//  [python]
def _get_data_protobuff(self, filename):
        filename = str(filename).split(',')
        filename_queue = tf.train.string_input_producer(filename,
                                                        num_epochs=None)
        reader = tf.TFRecordReader()
        _, serialized_example = reader.read(filename_queue)
        features = self._get_features(serialized_example)

        # image
        image, image_height, image_width = self._image_from_feature(features)

        # landmarks
        gt_heatmap, gt_lms, n_landmarks, mask_index, gt_mask = self._heatmap_from_feature(features)

        # infomations
        status = self._info_from_feature(features)

        # augmentation
        if self.augmentation:
            do_flip, do_rotate, do_scale = tf.unstack(self.augmentation_type())

            # rescale
            image_height = tf.to_int32(tf.to_float(image_height) * do_scale[0])
            image_width = tf.to_int32(tf.to_float(image_width) * do_scale[0])

            image = tf.image.resize_images(image, tf.stack([image_height, image_width]))
            gt_heatmap = tf.image.resize_images(gt_heatmap, tf.stack([image_height, image_width]))
            gt_lms = gt_lms*do_scale

            # rotate
            image = rotate_image_tensor(image, do_rotate)
            gt_heatmap = rotate_image_tensor(gt_heatmap, do_rotate)
            gt_lms = rotate_points_tensor(gt_lms, image, do_rotate)

        # crop to 256 * 256
        target_h = tf.to_int32(256)
        target_w = tf.to_int32(256)
        offset_h = tf.to_int32((image_height - target_h) / 2)
        offset_w = tf.to_int32((image_width - target_w) / 2)

        image = tf.image.crop_to_bounding_box(
            image, offset_h, offset_w, target_h, target_w)

        gt_heatmap = tf.image.crop_to_bounding_box(
            gt_heatmap, offset_h, offset_w, target_h, target_w)

        gt_lms = gt_lms - tf.to_float(tf.stack([offset_h, offset_w]))

        self._set_shape(image, gt_heatmap, gt_lms, mask_index, gt_mask)

        return image, gt_heatmap, gt_lms, mask_index, gt_mask
//  [python]
def _get_data_protobuff(self, filename):
        filename = str(filename).split(',')
        filename_queue = tf.train.string_input_producer(filename,
                                                        num_epochs=None)
        reader = tf.TFRecordReader()
        _, serialized_example = reader.read(filename_queue)
        features = self._get_features(serialized_example)

        # image
        image, image_height, image_width = self._image_from_feature(features)

        # landmarks
        gt_heatmap, gt_lms, n_landmarks, mask_index, gt_mask = self._heatmap_from_feature(features)

        # infomations
        status = self._info_from_feature(features)

        # augmentation
        if self.augmentation:
            do_flip, do_rotate, do_scale = tf.unstack(self.augmentation_type())

            # rescale
            image_height = tf.cast(tf.to_float(image_height) * do_scale[0], tf.int32)
            image_width = tf.cast(tf.to_float(image_width) * do_scale[0], tf.int32)

            image = tf.image.resize_images(image, tf.stack([image_height, image_width]))
            gt_heatmap = tf.image.resize_images(gt_heatmap, tf.stack([image_height, image_width]))
            gt_lms = gt_lms*do_scale

            # rotate
            image = rotate_image_tensor(image, do_rotate)
            gt_heatmap = rotate_image_tensor(gt_heatmap, do_rotate)
            gt_lms = rotate_points_tensor(gt_lms, image, do_rotate)

        # crop to 256 * 256
        target_h = tf.cast(256, tf.int32)
        target_w = tf.cast(256, tf.int32)
        offset_h = tf.cast((image_height - target_h) / 2, tf.int32)
        offset_w = tf.cast((image_width - target_w) / 2, tf.int32)

        image = tf.image.crop_to_bounding_box(
            image, offset_h, offset_w, target_h, target_w)

        gt_heatmap = tf.image.crop_to_bounding_box(
            gt_heatmap, offset_h, offset_w, target_h, target_w)

        gt_lms = gt_lms - tf.to_float(tf.stack([offset_h, offset_w]))

        self._set_shape(image, gt_heatmap, gt_lms, mask_index, gt_mask)

        return image, gt_heatmap, gt_lms, mask_index, gt_mask
//  [python]
def calc_accuracy_in_box(labels, coords, weights):
  """Calculate if coordinate lies within the ground truth.

  Args:
    labels: Ground truth label of shape [batch, height, width].
    coords: Regression output of shape [batch, 1, 1]
    weights: Padding mask for label

  Returns:
    Merged logits with shape [batch, height, width, num_classes].
  """
  labels *= weights
  labels = labels[:, :, :, 0]

  coords = tf.cast(coords, tf.int32)
  coords = tf.clip_by_value(coords, 0, FLAGS.image_size - 1)

  # Create the indices for the batch dimension
  batch_ind = tf.expand_dims(tf.range(tf.shape(coords)[0]), -1)

  # Concatenate the indices of the batch dimension with coords
  # to create a matrix of prediction indices
  coords = tf.cast(tf.concat([batch_ind, coords], -1), tf.int32)

  # Grab the elements prediced by index
  predicted_locations = tf.gather_nd(labels, coords)

  # Check if the elements selected are part of the ground truth bounding box
  accurate = tf.to_float(tf.greater(predicted_locations, 0))
  return accurate
//  [python]
def calc_accuracy_in_box(labels, coords, weights):
  """Calculate if coordinate lies within the ground truth.

  Args:
    labels: Ground truth label of shape [batch, height, width].
    coords: Regression output of shape [batch, 1, 1]
    weights: Padding mask for label

  Returns:
    Merged logits with shape [batch, height, width, num_classes].
  """
  labels *= weights
  labels = labels[:, :, :, 0]

  coords = tf.cast(coords, tf.int32)
  coords = tf.clip_by_value(coords, 0, FLAGS.image_size - 1)

  # Create the indices for the batch dimension
  batch_ind = tf.expand_dims(tf.range(tf.shape(coords)[0]), -1)

  # Concatenate the indices of the batch dimension with coords
  # to create a matrix of prediction indices
  coords = tf.cast(tf.concat([batch_ind, coords], -1), tf.int32)

  # Grab the elements prediced by index
  predicted_locations = tf.gather_nd(labels, coords)

  # Check if the elements selected are part of the ground truth bounding box
  accurate = tf.cast(tf.greater(predicted_locations, 0), tf.float32)
  return accurate
//  [python]
def __call__(self, image_height, image_width):
        """
        Note that we don't need to pass feature map shapes
        because we use only 'SAME' padding in all our networks.

        Arguments:
            image_height, image_width: scalar int tensors.
        Returns:
            a float tensor with shape [num_anchors, 4],
            boxes with normalized coordinates.
        """
        with tf.name_scope('anchor_generator'):

            image_height = tf.to_float(image_height)
            image_width = tf.to_float(image_width)

            feature_map_info = []
            num_anchors_per_feature_map = []
            for stride in self.strides:
                h = tf.to_int32(tf.ceil(image_height/stride))
                w = tf.to_int32(tf.ceil(image_width/stride))
                feature_map_info.append((stride, h, w))
                num_anchors_per_feature_map.append(h * w * self.num_anchors_per_location)

            # these are needed elsewhere
            self.num_anchors_per_feature_map = num_anchors_per_feature_map

            anchors = []

            # this is shared by all feature maps
            pairs = list(itertools.product(self.scale_multipliers, self.aspect_ratios))
            aspect_ratios = tf.constant([a for _, a in pairs], dtype=tf.float32)

            for i, (stride, h, w) in enumerate(feature_map_info):

                scales = tf.constant([m * self.scales[i] for m, _ in pairs], dtype=tf.float32)
                stride = tf.constant(stride, dtype=tf.float32)

                """
                It is true that
                image_height = h * stride - x, where 0 <= x < stride.

                Then image_height = (h - 1) * stride + (stride - x).
                So offset y must be equal to 0.5 * (stride - x).

                x = h * stride - image_height,
                y = 0.5 * (image_height - (h - 1) * stride),
                0 < y <= 0.5 * stride.

                Offset y is maximal when image_height is divisible by stride.
                Offset y is minimal when image_height = k * stride + 1, where k is a positive integer.
                """
                offset_y = 0.5 * (image_height - (tf.to_float(h) - 1.0) * stride)
                offset_x = 0.5 * (image_width - (tf.to_float(w) - 1.0) * stride)

                anchors.append(tile_anchors(
                    grid_height=h, grid_width=w,
                    scales=scales, aspect_ratios=aspect_ratios,
                    anchor_stride=(stride, stride),
                    anchor_offset=(offset_y, offset_x)
                ))

        with tf.name_scope('concatenate_normalize'):

            # this is for visualization and debugging only
            self.raw_anchors = anchors

            anchors = tf.concat(anchors, axis=0)

            # convert to the [0, 1] range
            scaler = tf.to_float(tf.stack([
                image_height, image_width,
                image_height, image_width
            ]))
            anchors /= scaler

        return anchors
//  [python]
def __call__(self, image_height, image_width):
        """
        Note that we don't need to pass feature map shapes
        because we use only 'SAME' padding in all our networks.

        Arguments:
            image_height, image_width: scalar int tensors.
        Returns:
            a float tensor with shape [num_anchors, 4],
            boxes with normalized coordinates.
        """
        with tf.name_scope('anchor_generator'):

            image_height = tf.cast(image_height, tf.float32)
            image_width = tf.cast(image_width, tf.float32)

            feature_map_info = []
            num_anchors_per_feature_map = []
            for stride in self.strides:
                h = tf.to_int32(tf.ceil(image_height/stride))
                w = tf.to_int32(tf.ceil(image_width/stride))
                feature_map_info.append((stride, h, w))
                num_anchors_per_feature_map.append(h * w * self.num_anchors_per_location)

            # these are needed elsewhere
            self.num_anchors_per_feature_map = num_anchors_per_feature_map

            anchors = []

            # this is shared by all feature maps
            pairs = list(itertools.product(self.scale_multipliers, self.aspect_ratios))
            aspect_ratios = tf.constant([a for _, a in pairs], dtype=tf.float32)

            for i, (stride, h, w) in enumerate(feature_map_info):

                scales = tf.constant([m * self.scales[i] for m, _ in pairs], dtype=tf.float32)
                stride = tf.constant(stride, dtype=tf.float32)

                """
                It is true that
                image_height = h * stride - x, where 0 <= x < stride.

                Then image_height = (h - 1) * stride + (stride - x).
                So offset y must be equal to 0.5 * (stride - x).

                x = h * stride - image_height,
                y = 0.5 * (image_height - (h - 1) * stride),
                0 < y <= 0.5 * stride.

                Offset y is maximal when image_height is divisible by stride.
                Offset y is minimal when image_height = k * stride + 1, where k is a positive integer.
                """
                offset_y = 0.5 * (image_height - (tf.to_float(h) - 1.0) * stride)
                offset_x = 0.5 * (image_width - (tf.to_float(w) - 1.0) * stride)

                anchors.append(tile_anchors(
                    grid_height=h, grid_width=w,
                    scales=scales, aspect_ratios=aspect_ratios,
                    anchor_stride=(stride, stride),
                    anchor_offset=(offset_y, offset_x)
                ))

        with tf.name_scope('concatenate_normalize'):

            # this is for visualization and debugging only
            self.raw_anchors = anchors

            anchors = tf.concat(anchors, axis=0)

            # convert to the [0, 1] range
            scaler = tf.to_float(tf.stack([
                image_height, image_width,
                image_height, image_width
            ]))
            anchors /= scaler

        return anchors
//  [python]
def tile_anchors(
        grid_height, grid_width,
        scales, aspect_ratios,
        anchor_stride, anchor_offset):
    """
    It returns boxes in absolute coordinates.

    Arguments:
        grid_height: a scalar int tensor, size of the grid in the y direction.
        grid_width: a scalar int tensor, size of the grid in the x direction.
        scales: a float tensor with shape [N],
            it represents the scale of each box in the basis set.
        aspect_ratios: a float tensor with shape [N],
            it represents the aspect ratio of each box in the basis set.
        anchor_stride: a tuple of float scalar tensors,
            difference in centers between anchors for adjacent grid positions.
        anchor_offset: a tuple of float scalar tensors,
            center of the anchor on upper left element of the grid ((0, 0)-th anchor).
    Returns:
        a float tensor with shape [grid_height * grid_width * N, 4].
    """
    N = tf.size(scales)
    ratio_sqrts = tf.sqrt(aspect_ratios)
    heights = scales / ratio_sqrts
    widths = scales * ratio_sqrts
    # widths/heights = aspect_ratios,
    # and scales = sqrt(heights * widths)

    # get a grid of box centers
    y_centers = tf.to_float(tf.range(grid_height)) * anchor_stride[0] + anchor_offset[0]
    x_centers = tf.to_float(tf.range(grid_width)) * anchor_stride[1] + anchor_offset[1]
    x_centers, y_centers = tf.meshgrid(x_centers, y_centers)
    # they have shape [grid_height, grid_width]

    centers = tf.stack([y_centers, x_centers], axis=2)
    centers = tf.expand_dims(centers, 2)
    centers = tf.tile(centers, [1, 1, N, 1])
    # shape [grid_height, grid_width, N, 2]

    sizes = tf.stack([heights, widths], axis=1)
    sizes = tf.expand_dims(tf.expand_dims(sizes, 0), 0)
    sizes = tf.tile(sizes, [grid_height, grid_width, 1, 1])
    # shape [grid_height, grid_width, N, 2]

    boxes = tf.concat([centers - 0.5 * sizes, centers + 0.5 * sizes], axis=3)
    # it has shape [grid_height, grid_width, N, 4]
    boxes = tf.reshape(boxes, [-1, 4])
    return boxes
//  [python]
def tile_anchors(
        grid_height, grid_width,
        scales, aspect_ratios,
        anchor_stride, anchor_offset):
    """
    It returns boxes in absolute coordinates.

    Arguments:
        grid_height: a scalar int tensor, size of the grid in the y direction.
        grid_width: a scalar int tensor, size of the grid in the x direction.
        scales: a float tensor with shape [N],
            it represents the scale of each box in the basis set.
        aspect_ratios: a float tensor with shape [N],
            it represents the aspect ratio of each box in the basis set.
        anchor_stride: a tuple of float scalar tensors,
            difference in centers between anchors for adjacent grid positions.
        anchor_offset: a tuple of float scalar tensors,
            center of the anchor on upper left element of the grid ((0, 0)-th anchor).
    Returns:
        a float tensor with shape [grid_height * grid_width * N, 4].
    """
    N = tf.size(scales)
    ratio_sqrts = tf.sqrt(aspect_ratios)
    heights = scales / ratio_sqrts
    widths = scales * ratio_sqrts
    # widths/heights = aspect_ratios,
    # and scales = sqrt(heights * widths)

    # get a grid of box centers
    y_centers = tf.cast(tf.range(grid_height), tf.float32) * anchor_stride[0] + anchor_offset[0]
    x_centers = tf.cast(tf.range(grid_width), tf.float32) * anchor_stride[1] + anchor_offset[1]
    x_centers, y_centers = tf.meshgrid(x_centers, y_centers)
    # they have shape [grid_height, grid_width]

    centers = tf.stack([y_centers, x_centers], axis=2)
    centers = tf.expand_dims(centers, 2)
    centers = tf.tile(centers, [1, 1, N, 1])
    # shape [grid_height, grid_width, N, 2]

    sizes = tf.stack([heights, widths], axis=1)
    sizes = tf.expand_dims(tf.expand_dims(sizes, 0), 0)
    sizes = tf.tile(sizes, [grid_height, grid_width, 1, 1])
    # shape [grid_height, grid_width, N, 2]

    boxes = tf.concat([centers - 0.5 * sizes, centers + 0.5 * sizes], axis=3)
    # it has shape [grid_height, grid_width, N, 4]
    boxes = tf.reshape(boxes, [-1, 4])
    return boxes
//  [python]
def decode_sequence(features, areas, hparams, decode_length,
                    post_processing=True):
  """Decodes the entire sequence in an auto-regressive way."""
  decode_utils.decode_n_step(seq2act_model.compute_logits,
                             features, areas,
                             hparams, n=decode_length, beam_size=1)
  if post_processing:
    features["input_refs"] = decode_utils.unify_input_ref(
        features["verbs"], features["input_refs"])
    pred_lengths = decode_utils.verb_refs_to_lengths(features["task"],
                                                     features["verb_refs"],
                                                     include_eos=False)
  predicted_actions = tf.concat([
      features["verb_refs"],
      features["obj_refs"],
      features["input_refs"],
      tf.to_int32(tf.expand_dims(features["verbs"], 2)),
      tf.to_int32(tf.expand_dims(features["objects"], 2))], axis=-1)
  if post_processing:
    predicted_actions = tf.where(
        tf.tile(tf.expand_dims(
            tf.sequence_mask(pred_lengths,
                             maxlen=tf.shape(predicted_actions)[1]),
            2), [1, 1, tf.shape(predicted_actions)[-1]]), predicted_actions,
        tf.zeros_like(predicted_actions))
  return predicted_actions
//  [python]
def decode_sequence(features, areas, hparams, decode_length,
                    post_processing=True):
  """Decodes the entire sequence in an auto-regressive way."""
  decode_utils.decode_n_step(seq2act_model.compute_logits,
                             features, areas,
                             hparams, n=decode_length, beam_size=1)
  if post_processing:
    features["input_refs"] = decode_utils.unify_input_ref(
        features["verbs"], features["input_refs"])
    pred_lengths = decode_utils.verb_refs_to_lengths(features["task"],
                                                     features["verb_refs"],
                                                     include_eos=False)
  predicted_actions = tf.concat([
      features["verb_refs"],
      features["obj_refs"],
      features["input_refs"],
      tf.cast(tf.expand_dims(features["verbs"], 2), tf.int32),
      tf.cast(tf.expand_dims(features["objects"], 2), tf.int32)], axis=-1)
  if post_processing:
    predicted_actions = tf.where(
        tf.tile(tf.expand_dims(
            tf.sequence_mask(pred_lengths,
                             maxlen=tf.shape(predicted_actions)[1]),
            2), [1, 1, tf.shape(predicted_actions)[-1]]), predicted_actions,
        tf.zeros_like(predicted_actions))
  return predicted_actions
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    data = self._dataset(num=self.test_size)

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.to_float(rs.rand(self.test_size, 28, 28, 1))
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.to_float(tf.to_float(image) > mask), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    data = self._dataset(num=self.test_size)

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.cast(rs.rand(self.test_size, 28, 28, 1), dtype=tf.float32)
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.cast(tf.cast(image, dtype=tf.float32) > mask, dtype=tf.float32), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def _get_arce_predictions(predictions, weights, label_weighted, labels,
                         up_weight, deprecated_rce,
                         total_positive, update_total_positive):
  """
  Returns the ARCE predictions, total_positive, update_total_positive and weights
  used by the rest of the twml.metrics.rce metric computation.
  """
  predictions_weighted = tf.multiply(predictions, weights, name="weighted_preds")
  label_weighted_comp = tf.subtract(tf.reduce_sum(weights), tf.reduce_sum(label_weighted))
  pred_weight_comp = tf.subtract(tf.reduce_sum(weights), tf.reduce_sum(predictions_weighted))
  normalizer_comp = label_weighted_comp / pred_weight_comp

  if up_weight is False:
    total_positive_unweighted = _metric_variable(
      name='total_positive_unweighted', shape=[], dtype=tf.float32)

    update_total_positive_unweighted = tf.assign_add(
      total_positive_unweighted, tf.reduce_sum(labels),
      name="total_positive_unweighted_update")

    if deprecated_rce:
      normalizer = tf.reduce_sum(labels) / tf.reduce_sum(label_weighted)
    else:
      # sum of labels / sum of weighted labels
      normalizer = update_total_positive_unweighted / update_total_positive

    label_comp = tf.subtract(tf.to_float(tf.size(labels)), tf.reduce_sum(labels))
    normalizer_comp = label_comp / label_weighted_comp

    # note that up_weight=True changes these for the rest of the twml.metric.rce computation
    weights = tf.ones(shape=tf.shape(labels), dtype=tf.float32, name="default_weight")
    total_positive = total_positive_unweighted
    update_total_positive = update_total_positive_unweighted
  else:
    if deprecated_rce:
      normalizer = tf.reduce_sum(label_weighted) / tf.reduce_sum(predictions_weighted)
    else:
      # normalizer used for NRCE (and ARCE with up_weight=True)
      total_prediction = _metric_variable(name='total_prediction', shape=[], dtype=tf.float32)

      # update the variable holding the sum of weighted predictions
      update_total_prediction = tf.assign_add(
        total_prediction, tf.reduce_sum(predictions_weighted), name="total_prediction_update")

      # this used to be tf.reduce_sum(label_weighted) / tf.reduce_sum(predictions_weighted)
      # but it measure normalizer over batch was too flawed an approximation.
      normalizer = update_total_positive / update_total_prediction

  pred_comp = tf.subtract(tf.ones(shape=tf.shape(labels), dtype=tf.float32), predictions)
  pred_comp_norm = tf.multiply(pred_comp, normalizer_comp, name="normalized_predictions_comp")
  pred_num = tf.multiply(predictions, normalizer, name="normalized_pred_numerator")
  pred_denom = tf.add(pred_num, pred_comp_norm, name="normalized_pred_denominator")
  predictions = pred_num / pred_denom

  return predictions, total_positive, update_total_positive, weights
//  [python]
def _get_arce_predictions(predictions, weights, label_weighted, labels,
                         up_weight, deprecated_rce,
                         total_positive, update_total_positive):
  """
  Returns the ARCE predictions, total_positive, update_total_positive and weights
  used by the rest of the twml.metrics.rce metric computation.
  """
  predictions_weighted = tf.multiply(predictions, weights, name="weighted_preds")
  label_weighted_comp = tf.subtract(tf.reduce_sum(weights), tf.reduce_sum(label_weighted))
  pred_weight_comp = tf.subtract(tf.reduce_sum(weights), tf.reduce_sum(predictions_weighted))
  normalizer_comp = label_weighted_comp / pred_weight_comp

  if up_weight is False:
    total_positive_unweighted = _metric_variable(
      name='total_positive_unweighted', shape=[], dtype=tf.float32)

    update_total_positive_unweighted = tf.assign_add(
      total_positive_unweighted, tf.reduce_sum(labels),
      name="total_positive_unweighted_update")

    if deprecated_rce:
      normalizer = tf.reduce_sum(labels) / tf.reduce_sum(label_weighted)
    else:
      # sum of labels / sum of weighted labels
      normalizer = update_total_positive_unweighted / update_total_positive

    label_comp = tf.subtract(tf.cast(tf.size(labels), tf.float32), tf.reduce_sum(labels))
    normalizer_comp = label_comp / label_weighted_comp

    # note that up_weight=True changes these for the rest of the twml.metric.rce computation
    weights = tf.ones(shape=tf.shape(labels), dtype=tf.float32, name="default_weight")
    total_positive = total_positive_unweighted
    update_total_positive = update_total_positive_unweighted
  else:
    if deprecated_rce:
      normalizer = tf.reduce_sum(label_weighted) / tf.reduce_sum(predictions_weighted)
    else:
      # normalizer used for NRCE (and ARCE with up_weight=True)
      total_prediction = _metric_variable(name='total_prediction', shape=[], dtype=tf.float32)

      # update the variable holding the sum of weighted predictions
      update_total_prediction = tf.assign_add(
        total_prediction, tf.reduce_sum(predictions_weighted), name="total_prediction_update")

      # this used to be tf.reduce_sum(label_weighted) / tf.reduce_sum(predictions_weighted)
      # but it measure normalizer over batch was too flawed an approximation.
      normalizer = update_total_positive / update_total_prediction

  pred_comp = tf.subtract(tf.ones(shape=tf.shape(labels), dtype=tf.float32), predictions)
  pred_comp_norm = tf.multiply(pred_comp, normalizer_comp, name="normalized_predictions_comp")
  pred_num = tf.multiply(predictions, normalizer, name="normalized_pred_numerator")
  pred_denom = tf.add(pred_num, pred_comp_norm, name="normalized_pred_denominator")
  predictions = pred_num / pred_denom

  return predictions, total_positive, update_total_positive, weights
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    if self._test_is_valid:
      data = self._dataset(
          os.path.join(self._data_dir, "train-images-idx3-ubyte"),
          os.path.join(self._data_dir, "train-labels-idx1-ubyte"))
      data = data.take(self.test_size)
    else:
      data = self._dataset(
          os.path.join(self._data_dir, "t10k-images-idx3-ubyte"),
          os.path.join(self._data_dir, "t10k-labels-idx1-ubyte"))

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.to_float(rs.rand(self.test_size, 28, 28, 1))
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.to_float(tf.to_float(image) > mask), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    if self._test_is_valid:
      data = self._dataset(
          os.path.join(self._data_dir, "train-images-idx3-ubyte"),
          os.path.join(self._data_dir, "train-labels-idx1-ubyte"))
      data = data.take(self.test_size)
    else:
      data = self._dataset(
          os.path.join(self._data_dir, "t10k-images-idx3-ubyte"),
          os.path.join(self._data_dir, "t10k-labels-idx1-ubyte"))

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.cast(rs.rand(self.test_size, 28, 28, 1), tf.float32)
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.cast(tf.cast(image, tf.float32) > mask, tf.float32), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def _import_feature(self, features, mesh, key):
    """Import a feature from the features dictionary into a mtf.Tensor.

    Args:
      features: a features dictionary
      mesh: a Mesh
      key: a string

    Returns:
      a mtf.Tensor with dtype int32 and shape self.batch_dims + self.length_dim
    """
    if key not in features:
      return None
    x = tf.to_int32(features[key])
    x = common_layers.expand_squeeze_to_nd(x, 2)
    batch_size = mtf.Shape(self.batch_dims).size
    x = x[:, :self.length_dim.size]
    extra_length = self.length_dim.size - tf.shape(x)[1]
    extra_batch = batch_size - tf.shape(x)[0]
    x = tf.pad(x, [[0, extra_batch], [0, extra_length]])
    mtf_shape = mtf.Shape(self.batch_dims + [self.length_dim])
    x = tf.reshape(x, mtf_shape.to_integer_list)
    return mtf.import_fully_replicated(mesh, x, mtf_shape, name=key)
//  [python]
def _import_feature(self, features, mesh, key):
    """Import a feature from the features dictionary into a mtf.Tensor.

    Args:
      features: a features dictionary
      mesh: a Mesh
      key: a string

    Returns:
      a mtf.Tensor with dtype int32 and shape self.batch_dims + self.length_dim
    """
    if key not in features:
      return None
    x = tf.cast(features[key], tf.int32)
    x = common_layers.expand_squeeze_to_nd(x, 2)
    batch_size = mtf.Shape(self.batch_dims).size
    x = x[:, :self.length_dim.size]
    extra_length = self.length_dim.size - tf.shape(x)[1]
    extra_batch = batch_size - tf.shape(x)[0]
    x = tf.pad(x, [[0, extra_batch], [0, extra_length]])
    mtf_shape = mtf.Shape(self.batch_dims + [self.length_dim])
    x = tf.reshape(x, mtf_shape.to_integer_list)
    return mtf.import_fully_replicated(mesh, x, mtf_shape, name=key)
//  [python]
def embedding_to_padding(emb):
  """Input embeddings -> is_padding."""
  emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1, keep_dims=True)
  return tf.to_float(tf.equal(emb_sum, 0.0))
//  [python]
def embedding_to_padding(emb):
  """Input embeddings -> is_padding."""
  emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1, keep_dims=True)
  return tf.cast(tf.equal(emb_sum, 0.0), tf.float32)
//  [python]
def create_model_fn(hparams, compute_additional_loss_fn=None,
                    compute_additional_metric_fn=None,
                    compute_seq_accuracy=False,
                    decode_length=20):
  """Creates the model function.

  Args:
    hparams: the hyper parameters.
    compute_additional_loss_fn: the optional callback for calculating
        additional loss.
    compute_additional_metric_fn: the optional callback for computing
        additional metrics.
    compute_seq_accuracy: whether to compute seq accuracy.
    decode_length: the maximum decoding length.
  Returns:
    the model function for estimator.
  """
  def model_fn(features, labels, mode):
    """The model function for creating an Estimtator."""
    del labels
    input_count = tf.reduce_sum(
        tf.to_int32(tf.greater(features["input_refs"][:, :, 1],
                               features["input_refs"][:, :, 0])))
    tf.summary.scalar("input_count", input_count)
    loss_dict, pred_dict, areas = seq2act_model.core_graph(
        features, hparams, mode, compute_additional_loss_fn)
    if mode == tf.estimator.ModeKeys.PREDICT:
      pred_dict["sequences"] = decode_sequence(
          features, areas, hparams, decode_length,
          post_processing=FLAGS.post_processing)
      return tf.estimator.EstimatorSpec(mode, predictions=pred_dict)
    elif mode == tf.estimator.ModeKeys.EVAL:
      metrics = {}
      _eval(metrics, pred_dict, loss_dict, features,
            areas, compute_seq_accuracy,
            hparams,
            metric_types=FLAGS.metric_types.split(","),
            decode_length=decode_length)
      if compute_additional_metric_fn:
        compute_additional_metric_fn(metrics, pred_dict, features)
      return tf.estimator.EstimatorSpec(
          mode, loss=loss_dict["total_loss"], eval_metric_ops=metrics)
    else:
      assert mode == tf.estimator.ModeKeys.TRAIN
      loss = loss_dict["total_loss"]
      for loss_name in loss_dict:
        if loss_name == "total_loss":
          continue
        if loss_name.endswith("losses"):
          continue
        tf.summary.scalar(loss_name, loss_dict[loss_name])
      step_num = tf.to_float(tf.train.get_global_step())
      schedule_string = hparams.learning_rate_schedule
      names = schedule_string.split("*")
      names = [name.strip() for name in names if name.strip()]
      ret = tf.constant(1.0)
      for name in names:
        ret *= learning_rate.learning_rate_factor(name, step_num, hparams)
      train_op = optimize.optimize(loss, ret, hparams)
      return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
  return model_fn
//  [python]
def create_model_fn(hparams, compute_additional_loss_fn=None,
                    compute_additional_metric_fn=None,
                    compute_seq_accuracy=False,
                    decode_length=20):
  """Creates the model function.

  Args:
    hparams: the hyper parameters.
    compute_additional_loss_fn: the optional callback for calculating
        additional loss.
    compute_additional_metric_fn: the optional callback for computing
        additional metrics.
    compute_seq_accuracy: whether to compute seq accuracy.
    decode_length: the maximum decoding length.
  Returns:
    the model function for estimator.
  """
  def model_fn(features, labels, mode):
    """The model function for creating an Estimtator."""
    del labels
    input_count = tf.reduce_sum(
        tf.cast(tf.greater(features["input_refs"][:, :, 1],
                               features["input_refs"][:, :, 0]), tf.int32))
    tf.summary.scalar("input_count", input_count)
    loss_dict, pred_dict, areas = seq2act_model.core_graph(
        features, hparams, mode, compute_additional_loss_fn)
    if mode == tf.estimator.ModeKeys.PREDICT:
      pred_dict["sequences"] = decode_sequence(
          features, areas, hparams, decode_length,
          post_processing=FLAGS.post_processing)
      return tf.estimator.EstimatorSpec(mode, predictions=pred_dict)
    elif mode == tf.estimator.ModeKeys.EVAL:
      metrics = {}
      _eval(metrics, pred_dict, loss_dict, features,
            areas, compute_seq_accuracy,
            hparams,
            metric_types=FLAGS.metric_types.split(","),
            decode_length=decode_length)
      if compute_additional_metric_fn:
        compute_additional_metric_fn(metrics, pred_dict, features)
      return tf.estimator.EstimatorSpec(
          mode, loss=loss_dict["total_loss"], eval_metric_ops=metrics)
    else:
      assert mode == tf.estimator.ModeKeys.TRAIN
      loss = loss_dict["total_loss"]
      for loss_name in loss_dict:
        if loss_name == "total_loss":
          continue
        if loss_name.endswith("losses"):
          continue
        tf.summary.scalar(loss_name, loss_dict[loss_name])
      step_num = tf.to_float(tf.train.get_global_step())
      schedule_string = hparams.learning_rate_schedule
      names = schedule_string.split("*")
      names = [name.strip() for name in names if name.strip()]
      ret = tf.constant(1.0)
      for name in names:
        ret *= learning_rate.learning_rate_factor(name, step_num, hparams)
      train_op = optimize.optimize(loss, ret, hparams)
      return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
  return model_fn
//  [python]
def get_splits(dataset, n_context, batch_size, points_perm=True):
  """Splits the dataset into target and context sets."""
  full_x, full_a, rewards = dataset
  full_y = rewards * full_a
  dataset_perm = np.random.permutation(len(full_x))[:batch_size]
  if points_perm:
    datapoints_perm = np.random.permutation(full_x.shape[1])
  else:
    datapoints_perm = np.arange(full_x.shape[1])

  target_x = tf.to_float(full_x[dataset_perm[:, None], datapoints_perm])
  target_y = tf.to_float(full_y[dataset_perm[:, None], datapoints_perm])
  target_a = tf.to_float(full_a[dataset_perm[:, None], datapoints_perm])
  context_x = target_x[:, :n_context, :]
  context_y = target_y[:, :n_context, :]
  unseen_target_y = target_y[:, n_context:]
  unseen_target_a = target_a[:, n_context:]
  return (context_x,
          context_y,
          target_x,
          target_y,
          unseen_target_y,
          unseen_target_a)
//  [python]
def get_splits(dataset, n_context, batch_size, points_perm=True):
  """Splits the dataset into target and context sets."""
  full_x, full_a, rewards = dataset
  full_y = rewards * full_a
  dataset_perm = np.random.permutation(len(full_x))[:batch_size]
  if points_perm:
    datapoints_perm = np.random.permutation(full_x.shape[1])
  else:
    datapoints_perm = np.arange(full_x.shape[1])

  target_x = tf.cast(full_x[dataset_perm[:, None], datapoints_perm], tf.float32)
  target_y = tf.cast(full_y[dataset_perm[:, None], datapoints_perm], tf.float32)
  target_a = tf.cast(full_a[dataset_perm[:, None], datapoints_perm], tf.float32)
  context_x = target_x[:, :n_context, :]
  context_y = target_y[:, :n_context, :]
  unseen_target_y = target_y[:, n_context:]
  unseen_target_a = target_a[:, n_context:]
  return (context_x,
          context_y,
          target_x,
          target_y,
          unseen_target_y,
          unseen_target_a)
//  [python]
def exp_term_for_constraints(rule_classes, num_classes, C):
	'''
	Func Desc:
	Compute the exponential term for the constraints

	Input:
	rule_classes ([num_rules,1]) - a list of classes associated with the rules
	num_classes (int)
	C

	Output:
	the required exponential term
	'''
	rule_classes_tensor = tf.to_float(tf.convert_to_tensor(rule_classes))
	#rule_classes_tensor = tf.reshape(rule_classes_tensor,[1,rule_classes])
	rule_classes_tensor = tf.expand_dims(rule_classes_tensor,0)
	class_types_tensor = tf.to_float(tf.convert_to_tensor(np.arange(num_classes).reshape(num_classes,1)))	
	#[num_classes,num_rules]
	class_rule_constraint = tf.to_float(tf.equal(class_types_tensor,rule_classes_tensor)) - 1.0
	class_rule_constraint = tf.exp(C*class_rule_constraint)
	return class_rule_constraint
//  [python]
def exp_term_for_constraints(rule_classes, num_classes, C):
	'''
	Func Desc:
	Compute the exponential term for the constraints

	Input:
	rule_classes ([num_rules,1]) - a list of classes associated with the rules
	num_classes (int)
	C

	Output:
	the required exponential term
	'''
	rule_classes_tensor = tf.cast(tf.convert_to_tensor(rule_classes), tf.float32)
	#rule_classes_tensor = tf.reshape(rule_classes_tensor,[1,rule_classes])
	rule_classes_tensor = tf.expand_dims(rule_classes_tensor,0)
	class_types_tensor = tf.cast(tf.convert_to_tensor(np.arange(num_classes).reshape(num_classes,1)), tf.float32)	
	#[num_classes,num_rules]
	class_rule_constraint = tf.cast(tf.equal(class_types_tensor,rule_classes_tensor), tf.float32) - 1.0
	class_rule_constraint = tf.exp(C*class_rule_constraint)
	return class_rule_constraint
//  [python]
def joint_scores_from_f_and_w(self,weights,m,f_probs):
		'''
		Func desc:
		Compute the learning scores obtained while jointly learning f(classification network) and w(rule network)

		Input:
		self object
		weights([num_instances, num_rules]) - the weights matrix corresponding to rule network(w_network) in the algorithm
		m([batch_size, num_rules]) - the rule coverage matrix where m_ij denotes if jth rule covers ith instance (if yes, then m_ij = 1)
		f_probs([batch_size, 1]) - the prob values from classification network (f_network)

		Output:
		results([batch_size,1]) -  
		'''
		num_classes = self.num_classes
		rule_classes = self.rule_classes
		#[batch_size, num_rules, 1]
		weights = tf.expand_dims(weights,-1)
		weights_mask = tf.to_float(tf.greater(weights,0.5))
		#[batch_size, num_rules, 1]
		m = tf.expand_dims(m,-1)
		m = m*weights_mask
		#[num_rules, num_classes]
		one_hot_rule_classes = tf.one_hot(rule_classes,num_classes,dtype=tf.float32)
		#[1, num_rules, num_classes]
		one_hot_rule_classes = tf.expand_dims(one_hot_rule_classes,0)
		#[batch_size, num_rules, num_classes]
		rule_weight_product = weights * one_hot_rule_classes + (1-weights)*(1-one_hot_rule_classes)
		sum_rule_firings = tf.reduce_sum(m,1)
		result = m*rule_weight_product #+ (1-m)
		#[batch_size, num_classes]
		result = tf.reduce_sum(result,1)/(sum_rule_firings+1e-20)
		result = result + f_probs
		return result
//  [python]
def joint_scores_from_f_and_w(self,weights,m,f_probs):
		'''
		Func desc:
		Compute the learning scores obtained while jointly learning f(classification network) and w(rule network)

		Input:
		self object
		weights([num_instances, num_rules]) - the weights matrix corresponding to rule network(w_network) in the algorithm
		m([batch_size, num_rules]) - the rule coverage matrix where m_ij denotes if jth rule covers ith instance (if yes, then m_ij = 1)
		f_probs([batch_size, 1]) - the prob values from classification network (f_network)

		Output:
		results([batch_size,1]) -  
		'''
		num_classes = self.num_classes
		rule_classes = self.rule_classes
		#[batch_size, num_rules, 1]
		weights = tf.expand_dims(weights,-1)
		weights_mask = tf.cast(tf.greater(weights,0.5),tf.float32)
		#[batch_size, num_rules, 1]
		m = tf.expand_dims(m,-1)
		m = m*weights_mask
		#[num_rules, num_classes]
		one_hot_rule_classes = tf.one_hot(rule_classes,num_classes,dtype=tf.float32)
		#[1, num_rules, num_classes]
		one_hot_rule_classes = tf.expand_dims(one_hot_rule_classes,0)
		#[batch_size, num_rules, num_classes]
		rule_weight_product = weights * one_hot_rule_classes + (1-weights)*(1-one_hot_rule_classes)
		sum_rule_firings = tf.reduce_sum(m,1)
		result = m*rule_weight_product #+ (1-m)
		#[batch_size, num_classes]
		result = tf.reduce_sum(result,1)/(sum_rule_firings+1e-20)
		result = result + f_probs
		return result
//  [python]
def compute_LL_phi(self, w_logits, weights, l, m, L, d, r):
		'''
		Func desc: 
		Computes the LL_phi term coming in the training objective

		Input:
		self object
		w_logits([batch_size, num_rules]) - 
		weights([batch_size, num_rules]) - the weights matrix corresponding to rule network(w_network) in the algorithm
		l([batch_size, num_rules]) - labels assigned by the rules
		m([batch_size, num_rules]) - the rule coverage matrix where m_ij = 1 if jth rule covers ith instance 
		L([batch_size, 1]) - L_i = 1 if the ith instance has already a label assigned to it in the dataset
		d([batch_size, 1]) - d_i = 1 if the ith instance is from labelled dataset
		r([batch_size, num_rules]) - the rule association matrix where r_ij = 1 if jth rule is associated with ith instance (r_ij = 1 => m_ij = 1)
		
		Output:
		loss(real number > 0) - the value of the LL_phi term
		'''
		psi = 1e-25
		L = tf.expand_dims(L,1)
		# [batch_size, num_rules]
		L = tf.tile(L,[1,self.num_rules])
		loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.to_float(tf.equal(l,L)),
													   logits=w_logits)
		loss = m*loss
		loss = (tf.to_float(tf.not_equal(l,L)) * loss) + (tf.to_float(tf.equal(l,L)) * r * loss)
# 		gcross_loss = gcross_utils.generalized_cross_entropy_bernoulli(weights,0.2)
		gcross_loss = generalized_cross_entropy_bernoulli(weights,0.2)
		gcross_loss = gcross_loss * m * tf.to_float(tf.equal(l,L)) * (1-r)
		loss = loss + gcross_loss
		loss = tf.reduce_sum(loss,axis=-1)
		loss = loss * d
		loss = tf.reduce_mean(loss)
		return loss
//  [python]
def compute_LL_phi(self, w_logits, weights, l, m, L, d, r):
		'''
		Func desc: 
		Computes the LL_phi term coming in the training objective

		Input:
		self object
		w_logits([batch_size, num_rules]) - 
		weights([batch_size, num_rules]) - the weights matrix corresponding to rule network(w_network) in the algorithm
		l([batch_size, num_rules]) - labels assigned by the rules
		m([batch_size, num_rules]) - the rule coverage matrix where m_ij = 1 if jth rule covers ith instance 
		L([batch_size, 1]) - L_i = 1 if the ith instance has already a label assigned to it in the dataset
		d([batch_size, 1]) - d_i = 1 if the ith instance is from labelled dataset
		r([batch_size, num_rules]) - the rule association matrix where r_ij = 1 if jth rule is associated with ith instance (r_ij = 1 => m_ij = 1)
		
		Output:
		loss(real number > 0) - the value of the LL_phi term
		'''
		psi = 1e-25
		L = tf.expand_dims(L,1)
		# [batch_size, num_rules]
		L = tf.tile(L,[1,self.num_rules])
		loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.cast(tf.equal(l,L),dtype=tf.float32),
													   logits=w_logits)
		loss = m*loss
		loss = (tf.cast(tf.not_equal(l,L),dtype=tf.float32) * loss) + (tf.cast(tf.equal(l,L),dtype=tf.float32) * r * loss)
# 		gcross_loss = gcross_utils.generalized_cross_entropy_bernoulli(weights,0.2)
		gcross_loss = generalized_cross_entropy_bernoulli(weights,0.2)
		gcross_loss = gcross_loss * m * tf.cast(tf.equal(l,L),dtype=tf.float32) * (1-r)
		loss = loss + gcross_loss
		loss = tf.reduce_sum(loss,axis=-1)
		loss = loss * d
		loss = tf.reduce_mean(loss)
		return loss
//  [python]
def _serving_input_receiver_fn():
  """Creates an input function for serving."""
  seq_len = FLAGS.max_seq_length
  serialized_example = tf.placeholder(
      dtype=tf.string, shape=[None], name="serialized_example")
  features = {
      "input_ids": tf.FixedLenFeature([seq_len], dtype=tf.int64),
      "input_mask": tf.FixedLenFeature([seq_len], dtype=tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_len], dtype=tf.int64),
  }
  feature_map = tf.parse_example(serialized_example, features=features)
  feature_map["is_real_example"] = tf.constant(1, dtype=tf.int32)
  feature_map["label_ids"] = tf.constant(0, dtype=tf.int32)

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

  return tf.estimator.export.ServingInputReceiver(
      features=feature_map, receiver_tensors=serialized_example)