jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def preprocess(self, inputs):
    frames, actions = inputs
    frames = tf.image.convert_image_dtype(frames, tf.float32)
    if self.is_discrete_action:
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.cast(actions, tf.float32)
    return frames, actions
//  [python]
def _convert_ids_to_strings(tgt_vocab_file, ids):
  """Convert prediction ids to words."""
  with tf.Session() as sess:
    reverse_target_vocab_table = lookup_ops.index_to_string_table_from_file(
        tgt_vocab_file, default_value=vocab_utils.UNK)
    sess.run(tf.tables_initializer())
    translations = sess.run(
        reverse_target_vocab_table.lookup(
            tf.to_int64(tf.convert_to_tensor(np.asarray(ids)))))
  return translations
//  [python]
def _convert_ids_to_strings(tgt_vocab_file, ids):
  """Convert prediction ids to words."""
  with tf.Session() as sess:
    reverse_target_vocab_table = lookup_ops.index_to_string_table_from_file(
        tgt_vocab_file, default_value=vocab_utils.UNK)
    sess.run(tf.tables_initializer())
    translations = sess.run(
        reverse_target_vocab_table.lookup(
            tf.cast(tf.convert_to_tensor(np.asarray(ids)), tf.int64)))
  return translations
//  [python]
def weak_cross_entropy_2d(y_pred, y_true, num_classes=None, epsilon=0.0001,
                          head=None):
    """ Weak Crossentropy 2d.

    Calculate the semantic segmentation using weak softmax cross entropy loss.

    Given the prediction `y_pred` shaped as 2d image and the corresponding
    y_true, this calculated the widely used semantic segmentation loss.
    Using `tf.nn.softmax_cross_entropy_with_logits` is currently not supported.
    See https://github.com/tensorflow/tensorflow/issues/2327#issuecomment-224491229

    Arguments:
        y_pred: `tensor, float` - [batch_size, width, height, num_classes].
        y_true: `Labels tensor, int32` - [batch_size, width, height, num_classes].
            The ground truth of your data.
        num_classes: `int`. Number of classes.
        epsilon: `float`. Small number to add to `y_pred`.
        head: `numpy array` - [num_classes]. Weighting the loss of each class.

    Returns:
        Loss tensor of type float.
    """
    if num_classes is None:
        num_classes = y_true.get_shape().as_list()[-1]
        # This only works if shape of y_true is defined
        assert (num_classes is not None)

    with tf.name_scope("weakCrossEntropy2d"):
        y_pred = tf.reshape(y_pred, (-1, num_classes))
        y_pred = y_pred + tf.constant(epsilon, dtype=y_pred.dtype)
        y_true = tf.to_float(tf.reshape(y_true, (-1, num_classes)))

        softmax = tf.nn.softmax(y_pred)

        if head is not None:
            cross_entropy = -tf.reduce_sum(tf.multiply(y_true * tf.log(softmax),
                                                  head), reduction_indices=[1])
        else:
            cross_entropy = -tf.reduce_sum(y_true * tf.log(softmax),
                                           reduction_indices=[1])

        cross_entropy_mean = tf.reduce_mean(cross_entropy,
                                            name="xentropy_mean")

    return cross_entropy_mean
//  [python]
def weak_cross_entropy_2d(y_pred, y_true, num_classes=None, epsilon=0.0001,
                          head=None):
    """ Weak Crossentropy 2d.

    Calculate the semantic segmentation using weak softmax cross entropy loss.

    Given the prediction `y_pred` shaped as 2d image and the corresponding
    y_true, this calculated the widely used semantic segmentation loss.
    Using `tf.nn.softmax_cross_entropy_with_logits` is currently not supported.
    See https://github.com/tensorflow/tensorflow/issues/2327#issuecomment-224491229

    Arguments:
        y_pred: `tensor, float` - [batch_size, width, height, num_classes].
        y_true: `Labels tensor, int32` - [batch_size, width, height, num_classes].
            The ground truth of your data.
        num_classes: `int`. Number of classes.
        epsilon: `float`. Small number to add to `y_pred`.
        head: `numpy array` - [num_classes]. Weighting the loss of each class.

    Returns:
        Loss tensor of type float.
    """
    if num_classes is None:
        num_classes = y_true.get_shape().as_list()[-1]
        # This only works if shape of y_true is defined
        assert (num_classes is not None)

    with tf.name_scope("weakCrossEntropy2d"):
        y_pred = tf.reshape(y_pred, (-1, num_classes))
        y_pred = y_pred + tf.constant(epsilon, dtype=y_pred.dtype)
        y_true = tf.cast(tf.reshape(y_true, (-1, num_classes)), dtype=tf.float32)

        softmax = tf.nn.softmax(y_pred)

        if head is not None:
            cross_entropy = -tf.reduce_sum(tf.multiply(y_true * tf.log(softmax),
                                                  head), reduction_indices=[1])
        else:
            cross_entropy = -tf.reduce_sum(y_true * tf.log(softmax),
                                           reduction_indices=[1])

        cross_entropy_mean = tf.reduce_mean(cross_entropy,
                                            name="xentropy_mean")

    return cross_entropy_mean
//  [python]
def compute_predictions(prediction_dict, references):
  """Predict the action tuple based on the logits."""
  prediction_dict["input"] = tf.argmax(references["input_logits"], -1)
  prediction_dict["verb_refs"] = predict_refs(references["verb_area_logits"],
                                              references["areas"]["starts"],
                                              references["areas"]["ends"])
  prediction_dict["obj_refs"] = predict_refs(references["obj_area_logits"],
                                             references["areas"]["starts"],
                                             references["areas"]["ends"])
  prediction_dict["input_refs"] = predict_refs(
      references["input_area_logits"],
      references["areas"]["starts"],
      references["areas"]["ends"]) * tf.to_int32(tf.expand_dims(
          prediction_dict["input"], 2))
//  [python]
def compute_predictions(prediction_dict, references):
  """Predict the action tuple based on the logits."""
  prediction_dict["input"] = tf.argmax(references["input_logits"], -1)
  prediction_dict["verb_refs"] = predict_refs(references["verb_area_logits"],
                                              references["areas"]["starts"],
                                              references["areas"]["ends"])
  prediction_dict["obj_refs"] = predict_refs(references["obj_area_logits"],
                                             references["areas"]["starts"],
                                             references["areas"]["ends"])
  prediction_dict["input_refs"] = predict_refs(
      references["input_area_logits"],
      references["areas"]["starts"],
      references["areas"]["ends"]) * tf.cast(tf.expand_dims(
          prediction_dict["input"], 2), tf.int32)
//  [python]
def compute_area_features(features, max_area_width, max_area_height=1, height=1,
                          epsilon=1e-6):
  """Computes features for each area.

  Args:
    features: a Tensor in a shape of [batch_size, height * width, depth].
    max_area_width: the max width allowed for an area.
    max_area_height: the max height allowed for an area.
    height: the height of the image.
    epsilon: the epsilon added to the variance for computing standard deviation.
  Returns:
    area_mean: A Tensor of shape [batch_size, num_areas, depth]
    area_std: A Tensor of shape [batch_size, num_areas, depth]
    area_sum: A Tensor of shape [batch_size, num_areas, depth]
    area_heights: A Tensor of shape [batch_size, num_areas, 1]
    area_widths: A Tensor of shape [batch_size, num_areas, 1]
  """
  with tf.name_scope("compute_area_features"):
    tf.logging.info("area_attention compute_area_features: %d x %d",
                    max_area_height, max_area_width)
    area_sum, area_heights, area_widths = _compute_sum_image(
        features, max_area_width=max_area_width,
        max_area_height=max_area_height, height=height)
    area_squared_sum, _, _ = _compute_sum_image(
        tf.pow(features, 2), max_area_width=max_area_width,
        max_area_height=max_area_height, height=height)
    sizes = tf.multiply(area_heights, area_widths)
    float_area_sizes = tf.to_float(sizes)
    area_mean = tf.div(area_sum, float_area_sizes)
    s2_n = tf.div(area_squared_sum, float_area_sizes)
    area_variance = tf.subtract(s2_n, tf.pow(area_mean, 2))
    area_std = tf.sqrt(tf.abs(area_variance) + epsilon)
    return area_mean, area_std, area_sum, area_heights, area_widths
//  [python]
def compute_area_features(features, max_area_width, max_area_height=1, height=1,
                          epsilon=1e-6):
  """Computes features for each area.

  Args:
    features: a Tensor in a shape of [batch_size, height * width, depth].
    max_area_width: the max width allowed for an area.
    max_area_height: the max height allowed for an area.
    height: the height of the image.
    epsilon: the epsilon added to the variance for computing standard deviation.
  Returns:
    area_mean: A Tensor of shape [batch_size, num_areas, depth]
    area_std: A Tensor of shape [batch_size, num_areas, depth]
    area_sum: A Tensor of shape [batch_size, num_areas, depth]
    area_heights: A Tensor of shape [batch_size, num_areas, 1]
    area_widths: A Tensor of shape [batch_size, num_areas, 1]
  """
  with tf.name_scope("compute_area_features"):
    tf.logging.info("area_attention compute_area_features: %d x %d",
                    max_area_height, max_area_width)
    area_sum, area_heights, area_widths = _compute_sum_image(
        features, max_area_width=max_area_width,
        max_area_height=max_area_height, height=height)
    area_squared_sum, _, _ = _compute_sum_image(
        tf.pow(features, 2), max_area_width=max_area_width,
        max_area_height=max_area_height, height=height)
    sizes = tf.multiply(area_heights, area_widths)
    float_area_sizes = tf.cast(sizes, tf.float32)
    area_mean = tf.div(area_sum, float_area_sizes)
    s2_n = tf.div(area_squared_sum, float_area_sizes)
    area_variance = tf.subtract(s2_n, tf.pow(area_mean, 2))
    area_std = tf.sqrt(tf.abs(area_variance) + epsilon)
    return area_mean, area_std, area_sum, area_heights, area_widths
//  [python]
def batch_norm_with_mask(x, is_training, mask, num_channels, name="bn",
                         decay=0.9, epsilon=1e-3, data_format="NHWC"):

  shape = [num_channels]
  indices = tf.where(mask)
  indices = tf.to_int32(indices)
  indices = tf.reshape(indices, [-1])

  with tf.variable_scope(name, reuse=None if is_training else True):
    offset = tf.get_variable(
      "offset", shape,
      initializer=tf.constant_initializer(0.0, dtype=tf.float32))
    scale = tf.get_variable(
      "scale", shape,
      initializer=tf.constant_initializer(1.0, dtype=tf.float32))
    offset = tf.boolean_mask(offset, mask)
    scale = tf.boolean_mask(scale, mask)

    moving_mean = tf.get_variable(
      "moving_mean", shape, trainable=False,
      initializer=tf.constant_initializer(0.0, dtype=tf.float32))
    moving_variance = tf.get_variable(
      "moving_variance", shape, trainable=False,
      initializer=tf.constant_initializer(1.0, dtype=tf.float32))

    if is_training:
      x, mean, variance = tf.nn.fused_batch_norm(
        x, scale, offset, epsilon=epsilon, data_format=data_format,
        is_training=True)
      mean = (1.0 - decay) * (tf.boolean_mask(moving_mean, mask) - mean)
      variance = (1.0 - decay) * (tf.boolean_mask(moving_variance, mask) - variance)
      update_mean = tf.scatter_sub(moving_mean, indices, mean, use_locking=True)
      update_variance = tf.scatter_sub(
        moving_variance, indices, variance, use_locking=True)
      with tf.control_dependencies([update_mean, update_variance]):
        x = tf.identity(x)
    else:
      masked_moving_mean = tf.boolean_mask(moving_mean, mask)
      masked_moving_variance = tf.boolean_mask(moving_variance, mask)
      x, _, _ = tf.nn.fused_batch_norm(x, scale, offset,
                                       mean=masked_moving_mean,
                                       variance=masked_moving_variance,
                                       epsilon=epsilon, data_format=data_format,
                                       is_training=False)
  return x
//  [python]
def batch_norm_with_mask(x, is_training, mask, num_channels, name="bn",
                         decay=0.9, epsilon=1e-3, data_format="NHWC"):

  shape = [num_channels]
  indices = tf.where(mask)
  indices = tf.cast(indices, tf.int32)
  indices = tf.reshape(indices, [-1])

  with tf.variable_scope(name, reuse=None if is_training else True):
    offset = tf.get_variable(
      "offset", shape,
      initializer=tf.constant_initializer(0.0, dtype=tf.float32))
    scale = tf.get_variable(
      "scale", shape,
      initializer=tf.constant_initializer(1.0, dtype=tf.float32))
    offset = tf.boolean_mask(offset, mask)
    scale = tf.boolean_mask(scale, mask)

    moving_mean = tf.get_variable(
      "moving_mean", shape, trainable=False,
      initializer=tf.constant_initializer(0.0, dtype=tf.float32))
    moving_variance = tf.get_variable(
      "moving_variance", shape, trainable=False,
      initializer=tf.constant_initializer(1.0, dtype=tf.float32))

    if is_training:
      x, mean, variance = tf.nn.fused_batch_norm(
        x, scale, offset, epsilon=epsilon, data_format=data_format,
        is_training=True)
      mean = (1.0 - decay) * (tf.boolean_mask(moving_mean, mask) - mean)
      variance = (1.0 - decay) * (tf.boolean_mask(moving_variance, mask) - variance)
      update_mean = tf.scatter_sub(moving_mean, indices, mean, use_locking=True)
      update_variance = tf.scatter_sub(
        moving_variance, indices, variance, use_locking=True)
      with tf.control_dependencies([update_mean, update_variance]):
        x = tf.identity(x)
    else:
      masked_moving_mean = tf.boolean_mask(moving_mean, mask)
      masked_moving_variance = tf.boolean_mask(moving_variance, mask)
      x, _, _ = tf.nn.fused_batch_norm(x, scale, offset,
                                       mean=masked_moving_mean,
                                       variance=masked_moving_variance,
                                       epsilon=epsilon, data_format=data_format,
                                       is_training=False)
  return x
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']
        seqs_ts = features['seqs_t'] / self.time_scale

        # Embedding and Transform
        seqs_units = self.item_embs(seqs_id)
        # time delta hereby is sensitive to the timepoints of purchasing
        seqs_spans = tf.tile(
            tf.expand_dims(seqs_ts[:, 1:], 2), [1, 1, self.seqslen]) - tf.tile(
            tf.expand_dims(seqs_ts[:, :-1], 1), [1, self.seqslen, 1])
        seqs_spans = tf.to_float(tf.maximum(seqs_spans, 0.))

        # Dropout
        seqs_units = tf.layers.dropout(seqs_units, rate=self.hidden_dropout_rate,
                                       training=tf.convert_to_tensor(is_training))
        seqs_masks = tf.expand_dims(tf.to_float(tf.not_equal(seqs_id, 0)), -1)

        # multi-head attention
        seqs_outs = seqs_units * seqs_masks
        for i, (attention, dense) in enumerate(zip(self.list_attention, self.list_dense)):
            with tf.variable_scope("num_blocks_%d" % i):
                with tf.variable_scope("attention"):
                    seqs_outs = attention(layernorm(seqs_outs), seqs_outs, seqs_spans, is_training, causality=True)

                with tf.variable_scope("feedforward"):
                    seqs_outs = dense(layernorm(seqs_outs), is_training)
                    seqs_outs *= seqs_masks

        with tf.variable_scope("out_ln"):
            seqs_outs = layernorm(seqs_outs)
        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']
        seqs_ts = features['seqs_t'] / self.time_scale

        # Embedding and Transform
        seqs_units = self.item_embs(seqs_id)
        # time delta hereby is sensitive to the timepoints of purchasing
        seqs_spans = tf.tile(
            tf.expand_dims(seqs_ts[:, 1:], 2), [1, 1, self.seqslen]) - tf.tile(
            tf.expand_dims(seqs_ts[:, :-1], 1), [1, self.seqslen, 1])
        seqs_spans = tf.cast(tf.maximum(seqs_spans, 0.), dtype=tf.float32)

        # Dropout
        seqs_units = tf.layers.dropout(seqs_units, rate=self.hidden_dropout_rate,
                                       training=tf.convert_to_tensor(is_training))
        seqs_masks = tf.expand_dims(tf.cast(tf.not_equal(seqs_id, 0), dtype=tf.float32), -1)

        # multi-head attention
        seqs_outs = seqs_units * seqs_masks
        for i, (attention, dense) in enumerate(zip(self.list_attention, self.list_dense)):
            with tf.variable_scope("num_blocks_%d" % i):
                with tf.variable_scope("attention"):
                    seqs_outs = attention(layernorm(seqs_outs), seqs_outs, seqs_spans, is_training, causality=True)

                with tf.variable_scope("feedforward"):
                    seqs_outs = dense(layernorm(seqs_outs), is_training)
                    seqs_outs *= seqs_masks

        with tf.variable_scope("out_ln"):
            seqs_outs = layernorm(seqs_outs)
        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits
//  [python]
def accuracy(label, logits):
  """Computes accuracy from given label and logits."""
  return tf.reduce_mean(tf.to_float(tf.equal(label, tf.argmax(logits, axis=1))))
//  [python]
def accuracy(label, logits):
  """Computes accuracy from given label and logits."""
  return tf.reduce_mean(tf.cast(tf.equal(label, tf.argmax(logits, axis=1)), tf.float32))
//  [python]
def compute_losses(loss_dict, features, references, hparams):
  """Compute the loss based on the logits and labels."""
  # Commented code can be useful for examining seq lengths distribution
  # srcs, _, counts = tf.unique_with_counts(features["data_source"])
  # lengths = tf.reduce_sum(tf.to_int32(
  #    tf.greater(features["verb_refs"][:, :, 1],
  #               features["verb_refs"][:, :, 0])), -1) - 1
  # lengths, _, len_counts = tf.unique_with_counts(lengths)
  # with tf.control_dependencies([
  #    tf.print("sources", srcs, counts, lengths, len_counts, summarize=1000)]):
  input_mask = tf.to_float(
      tf.greater(features["verb_refs"][:, :, 1],
                 features["verb_refs"][:, :, 0]))
  input_loss = tf.reduce_mean(
      tf.losses.sparse_softmax_cross_entropy(
          labels=tf.to_int32(
              tf.greater(features["input_refs"][:, :, 1],
                         features["input_refs"][:, :, 0])),
          logits=references["input_logits"],
          reduction=tf.losses.Reduction.NONE) * input_mask)
  encoder_input_length = common_layers.shape_list(features["task"])[1]
  verb_area_loss = area_utils.area_loss(
      logits=references["verb_area_logits"], ranges=features["verb_refs"],
      length=encoder_input_length,
      max_area_width=hparams.max_span)
  object_area_loss = area_utils.area_loss(
      logits=references["obj_area_logits"],
      ranges=features["obj_refs"],
      length=encoder_input_length,
      max_area_width=hparams.max_span)
  input_area_loss = area_utils.area_loss(
      logits=references["input_area_logits"], ranges=features["input_refs"],
      length=encoder_input_length,
      max_area_width=hparams.max_span)
  loss_dict["reference_loss"] = (
      input_loss + verb_area_loss + object_area_loss + input_area_loss)
  loss_dict["input_loss"] = input_loss
  loss_dict["verb_refs_loss"] = verb_area_loss
  loss_dict["obj_refs_loss"] = object_area_loss
  loss_dict["input_refs_loss"] = input_area_loss
  return loss_dict["reference_loss"]
//  [python]
def compute_losses(loss_dict, features, references, hparams):
  """Compute the loss based on the logits and labels."""
  # Commented code can be useful for examining seq lengths distribution
  # srcs, _, counts = tf.unique_with_counts(features["data_source"])
  # lengths = tf.reduce_sum(tf.to_int32(
  #    tf.greater(features["verb_refs"][:, :, 1],
  #               features["verb_refs"][:, :, 0])), -1) - 1
  # lengths, _, len_counts = tf.unique_with_counts(lengths)
  # with tf.control_dependencies([
  #    tf.print("sources", srcs, counts, lengths, len_counts, summarize=1000)]):
  input_mask = tf.cast(
      tf.greater(features["verb_refs"][:, :, 1],
                 features["verb_refs"][:, :, 0]), tf.float32)
  input_loss = tf.reduce_mean(
      tf.losses.sparse_softmax_cross_entropy(
          labels=tf.to_int32(
              tf.greater(features["input_refs"][:, :, 1],
                         features["input_refs"][:, :, 0])),
          logits=references["input_logits"],
          reduction=tf.losses.Reduction.NONE) * input_mask)
  encoder_input_length = common_layers.shape_list(features["task"])[1]
  verb_area_loss = area_utils.area_loss(
      logits=references["verb_area_logits"], ranges=features["verb_refs"],
      length=encoder_input_length,
      max_area_width=hparams.max_span)
  object_area_loss = area_utils.area_loss(
      logits=references["obj_area_logits"],
      ranges=features["obj_refs"],
      length=encoder_input_length,
      max_area_width=hparams.max_span)
  input_area_loss = area_utils.area_loss(
      logits=references["input_area_logits"], ranges=features["input_refs"],
      length=encoder_input_length,
      max_area_width=hparams.max_span)
  loss_dict["reference_loss"] = (
      input_loss + verb_area_loss + object_area_loss + input_area_loss)
  loss_dict["input_loss"] = input_loss
  loss_dict["verb_refs_loss"] = verb_area_loss
  loss_dict["obj_refs_loss"] = object_area_loss
  loss_dict["input_refs_loss"] = input_area_loss
  return loss_dict["reference_loss"]
//  [python]
def _build_loss(self):
        """Build our cross entropy loss."""

        with tf.variable_scope("Loss", reuse=tf.AUTO_REUSE):

            # sh = tf.shape(self.x_in_2)
            x1_ = tf.squeeze(self.x_in_1, axis=1)
            x2_ = tf.squeeze(self.x_in_2, axis=1)


            print(self.R_hat.shape)
            self.x2_hat = tf_matrix_vector_mul(self.R_hat, x1_)
            self.x2_hat = tf_add_vectors(self.x2_hat, self.t_hat)
            sub = self.x2_hat - x2_

            if self.config.loss_function == 'wls':
                w_mul = tf.expand_dims(self.weights, axis=2)
                w_mul1 = tf.tile(w_mul, [1, 1, 3])
                sub = w_mul1*sub

            print(sub.shape)
            r_loss = self.loss_function(sub)
            r_loss = tf.reduce_sum(r_loss, axis=1)

            self.rec_loss = tf.reduce_mean(r_loss)
            self.rec_loss =(self.config.loss_reconstruction * self.rec_loss * tf.to_float(
                            self.global_step >= tf.to_int64(
                            self.config.loss_reconstruction_init_iter)))

            tf.summary.scalar("reconstruction_loss", self.rec_loss)

            is_pos = tf.to_float(self.f_in > 0)
            is_neg = tf.to_float(self.f_in <= 0)
            c = is_pos - is_neg

            # clss
            clf_losses = -tf.log(tf.nn.sigmoid(c * self.logits))

            num_pos = tf.nn.relu(tf.reduce_sum(is_pos, axis=1) - 1.0) + 1.0
            num_neg = tf.nn.relu(tf.reduce_sum(is_neg, axis=1) - 1.0) + 1.0
            classif_loss_p = tf.reduce_sum(
                clf_losses * is_pos, axis=1
            )
            classif_loss_n = tf.reduce_sum(
                clf_losses * is_neg, axis=1
            )
            self.clf_loss = tf.reduce_mean(
                classif_loss_p * 0.5 / num_pos +
                classif_loss_n * 0.5 / num_neg
            )

            self.clf_loss = self.config.loss_classif * self.clf_loss
            tf.summary.scalar("classification_loss", self.clf_loss)
            tf.summary.scalar(
                "classif_loss_p",
                tf.reduce_mean(classif_loss_p * 0.5 / num_pos))
            tf.summary.scalar(
                "classif_loss_n",
                tf.reduce_mean(classif_loss_n * 0.5 / num_neg))

            # L2 loss

            for var in tf.trainable_variables():
                if "weights" in var.name:
                    tf.add_to_collection("l2_losses", tf.reduce_sum(var ** 2))
            l2_loss = tf.add_n(tf.get_collection("l2_losses"))
            tf.summary.scalar("l2_loss", l2_loss)

            self.loss = self.config.loss_decay * l2_loss

            self.loss += self.rec_loss
            self.loss += self.clf_loss

            tf.summary.scalar("loss", self.loss)
//  [python]
def _build_loss(self):
        """Build our cross entropy loss."""

        with tf.variable_scope("Loss", reuse=tf.AUTO_REUSE):

            # sh = tf.shape(self.x_in_2)
            x1_ = tf.squeeze(self.x_in_1, axis=1)
            x2_ = tf.squeeze(self.x_in_2, axis=1)


            print(self.R_hat.shape)
            self.x2_hat = tf_matrix_vector_mul(self.R_hat, x1_)
            self.x2_hat = tf_add_vectors(self.x2_hat, self.t_hat)
            sub = self.x2_hat - x2_

            if self.config.loss_function == 'wls':
                w_mul = tf.expand_dims(self.weights, axis=2)
                w_mul1 = tf.tile(w_mul, [1, 1, 3])
                sub = w_mul1*sub

            print(sub.shape)
            r_loss = self.loss_function(sub)
            r_loss = tf.reduce_sum(r_loss, axis=1)

            self.rec_loss = tf.reduce_mean(r_loss)
            self.rec_loss =(self.config.loss_reconstruction * self.rec_loss * tf.cast(
                            self.global_step >= tf.to_int64(
                            self.config.loss_reconstruction_init_iter), tf.float32))

            tf.summary.scalar("reconstruction_loss", self.rec_loss)

            is_pos = tf.cast(self.f_in > 0, tf.float32)
            is_neg = tf.cast(self.f_in <= 0, tf.float32)
            c = is_pos - is_neg

            # clss
            clf_losses = -tf.log(tf.nn.sigmoid(c * self.logits))

            num_pos = tf.nn.relu(tf.reduce_sum(is_pos, axis=1) - 1.0) + 1.0
            num_neg = tf.nn.relu(tf.reduce_sum(is_neg, axis=1) - 1.0) + 1.0
            classif_loss_p = tf.reduce_sum(
                clf_losses * is_pos, axis=1
            )
            classif_loss_n = tf.reduce_sum(
                clf_losses * is_neg, axis=1
            )
            self.clf_loss = tf.reduce_mean(
                classif_loss_p * 0.5 / num_pos +
                classif_loss_n * 0.5 / num_neg
            )

            self.clf_loss = self.config.loss_classif * self.clf_loss
            tf.summary.scalar("classification_loss", self.clf_loss)
            tf.summary.scalar(
                "classif_loss_p",
                tf.reduce_mean(classif_loss_p * 0.5 / num_pos))
            tf.summary.scalar(
                "classif_loss_n",
                tf.reduce_mean(classif_loss_n * 0.5 / num_neg))

            # L2 loss

            for var in tf.trainable_variables():
                if "weights" in var.name:
                    tf.add_to_collection("l2_losses", tf.reduce_sum(var ** 2))
            l2_loss = tf.add_n(tf.get_collection("l2_losses"))
            tf.summary.scalar("l2_loss", l2_loss)

            self.loss = self.config.loss_decay * l2_loss

            self.loss += self.rec_loss
            self.loss += self.clf_loss

            tf.summary.scalar("loss", self.loss)
//  [python]
def log_prob(self, x_final, num_samples=1):
    """Compute log probability lower bound on x_final.

    Args:
      x_final: [batch_size] + data_dim tensor.
      num_samples: Optional number of samples to compute bounds.
    Returns:
      log probability lower bound.
    """
    tiled_x_final = tf.tile(x_final, [num_samples] + [1] * len(self.data_dim))
    q = self.q(tiled_x_final)
    rho_final = q.sample()  # [num_samples * batch_size, data_dim]

    x_0, rho_0 = self._reverse_hamiltonian_dynamics(tiled_x_final, rho_final)
    elbo = (
        self.proposal.log_prob(x_0) + self.momentum_proposal.log_prob(rho_0) -
        q.log_prob(rho_final))
    iwae = (tf.reduce_logsumexp(tf.reshape(elbo, [num_samples, -1]), axis=0)
            - tf.log(tf.to_float(num_samples)))
    return iwae
//  [python]
def log_prob(self, x_final, num_samples=1):
    """Compute log probability lower bound on x_final.

    Args:
      x_final: [batch_size] + data_dim tensor.
      num_samples: Optional number of samples to compute bounds.
    Returns:
      log probability lower bound.
    """
    tiled_x_final = tf.tile(x_final, [num_samples] + [1] * len(self.data_dim))
    q = self.q(tiled_x_final)
    rho_final = q.sample()  # [num_samples * batch_size, data_dim]

    x_0, rho_0 = self._reverse_hamiltonian_dynamics(tiled_x_final, rho_final)
    elbo = (
        self.proposal.log_prob(x_0) + self.momentum_proposal.log_prob(rho_0) -
        q.log_prob(rho_final))
    iwae = (tf.reduce_logsumexp(tf.reshape(elbo, [num_samples, -1]), axis=0)
            - tf.log(tf.cast(num_samples, tf.float32)))
    return iwae
//  [python]
def _torso(self, input_):
    """Processing of all the visual and language inputs to the LSTM core."""

    # Extract the inputs
    last_action, env_output = input_
    last_reward, _, _, observation = env_output
    frame = observation[self._idx_frame]
    goal = observation[self._idx_goal]
    goal = tf.to_float(goal)

    # Convert to image to floats and normalise.
    frame = tf.to_float(frame)
    frame = snt.FlattenTrailingDimensions(dim_from=3)(frame)
    frame /= 255.0

    # Feed image through convnet.
    with tf.variable_scope('convnet'):
      # Convolutional layers.
      conv_out = self._convnet(frame)
      # Fully connected layer.
      conv_out = snt.BatchFlatten()(conv_out)
      conv_out = snt.Linear(256)(conv_out)
      conv_out = tf.nn.relu(conv_out)

    # Concatenate outputs of the visual and instruction pathways.
    if self._feed_action_and_reward:
      # Append clipped last reward and one hot last action.
      tf.logging.info('Append last reward clipped to: %f', self._max_reward)
      clipped_last_reward = tf.expand_dims(
          tf.clip_by_value(last_reward, -self._max_reward, self._max_reward),
          -1)
      tf.logging.info('Append last action (one-hot of %d)', self._num_actions)
      one_hot_last_action = tf.one_hot(last_action, self._num_actions)
      tf.logging.info('Append goal:')
      tf.logging.info(goal)
      action_and_reward = tf.concat([clipped_last_reward, one_hot_last_action],
                                    axis=1)
    else:
      action_and_reward = tf.constant([0], dtype=tf.float32)
    return conv_out, action_and_reward, goal
//  [python]
def _torso(self, input_):
    """Processing of all the visual and language inputs to the LSTM core."""

    # Extract the inputs
    last_action, env_output = input_
    last_reward, _, _, observation = env_output
    frame = observation[self._idx_frame]
    goal = observation[self._idx_goal]
    goal = tf.cast(goal, tf.float32)

    # Convert to image to floats and normalise.
    frame = tf.cast(frame, tf.float32)
    frame = snt.FlattenTrailingDimensions(dim_from=3)(frame)
    frame /= 255.0

    # Feed image through convnet.
    with tf.variable_scope('convnet'):
      # Convolutional layers.
      conv_out = self._convnet(frame)
      # Fully connected layer.
      conv_out = snt.BatchFlatten()(conv_out)
      conv_out = snt.Linear(256)(conv_out)
      conv_out = tf.nn.relu(conv_out)

    # Concatenate outputs of the visual and instruction pathways.
    if self._feed_action_and_reward:
      # Append clipped last reward and one hot last action.
      tf.logging.info('Append last reward clipped to: %f', self._max_reward)
      clipped_last_reward = tf.expand_dims(
          tf.clip_by_value(last_reward, -self._max_reward, self._max_reward),
          -1)
      tf.logging.info('Append last action (one-hot of %d)', self._num_actions)
      one_hot_last_action = tf.one_hot(last_action, self._num_actions)
      tf.logging.info('Append goal:')
      tf.logging.info(goal)
      action_and_reward = tf.concat([clipped_last_reward, one_hot_last_action],
                                    axis=1)
    else:
      action_and_reward = tf.constant([0], dtype=tf.float32)
    return conv_out, action_and_reward, goal
//  [python]
def __init__(self, anchors, max_grid, batch_size, warmup_batches, ignore_thresh, 
                    grid_scale, obj_scale, noobj_scale, xywh_scale, class_scale, 
                    **kwargs):
        # make the model settings persistent
        self.ignore_thresh  = ignore_thresh
        self.warmup_batches = warmup_batches
        self.anchors        = tf.constant(anchors, dtype='float', shape=[1,1,1,3,2])
        self.grid_scale     = grid_scale
        self.obj_scale      = obj_scale
        self.noobj_scale    = noobj_scale
        self.xywh_scale     = xywh_scale
        self.class_scale    = class_scale        

        # make a persistent mesh grid
        max_grid_h, max_grid_w = max_grid

        cell_x = tf.to_float(tf.reshape(tf.tile(tf.range(max_grid_w), [max_grid_h]), (1, max_grid_h, max_grid_w, 1, 1)))
        cell_y = tf.transpose(cell_x, (0,2,1,3,4))
        self.cell_grid = tf.tile(tf.concat([cell_x,cell_y],-1), [batch_size, 1, 1, 3, 1])

        super(YoloLayer, self).__init__(**kwargs)
//  [python]
def __init__(self, anchors, max_grid, batch_size, warmup_batches, ignore_thresh, 
                    grid_scale, obj_scale, noobj_scale, xywh_scale, class_scale, 
                    **kwargs):
        # make the model settings persistent
        self.ignore_thresh  = ignore_thresh
        self.warmup_batches = warmup_batches
        self.anchors        = tf.constant(anchors, dtype='float', shape=[1,1,1,3,2])
        self.grid_scale     = grid_scale
        self.obj_scale      = obj_scale
        self.noobj_scale    = noobj_scale
        self.xywh_scale     = xywh_scale
        self.class_scale    = class_scale        

        # make a persistent mesh grid
        max_grid_h, max_grid_w = max_grid

        cell_x = tf.cast(tf.reshape(tf.tile(tf.range(max_grid_w), [max_grid_h]), (1, max_grid_h, max_grid_w, 1, 1)),dtype=tf.float32)
        cell_y = tf.transpose(cell_x, (0,2,1,3,4))
        self.cell_grid = tf.tile(tf.concat([cell_x,cell_y],-1), [batch_size, 1, 1, 3, 1])

        super(YoloLayer, self).__init__(**kwargs)
//  [python]
def mask_logits(vec, mask):
  """Mask `vec` in log-space.

  Elements in `vec` that are not in `mask` will be set to be very
  negative, so that when no longer in log-space (e.i., after `tf.exp(vec)`)
  their values will be very close to zero.

  Args:
    vec: <float32>[...] tensor to mask
    mask: Either a None (in which case this is a no-op), a boolean or 0/1 float
      mask that matches all, or all but the last, dimensions of `vec`, or 1-D
      integer length mask

  Raises:
    ValueError: If `mask` cannot be matched to `vec`
  Returns:
    masked: vec:<float32>[...]
  """
  if mask is None:
    return vec

  if mask.dtype == tf.int32:
    # Assume `mask` holds sequence lengths
    if len(vec.shape) not in [2, 3]:
      raise ValueError("Can't use a length mask on tensor of rank>3")
    mask = tf.sequence_mask(mask, tf.shape(vec)[1], tf.float32)
  else:
    mask = tf.to_float(mask)

  if len(mask.shape) == (len(vec.shape) - 1):
    mask = tf.expand_dims(mask, len(vec.shape) - 1)

  return vec * mask - (1 - mask) * 1E20
//  [python]
def mask_logits(vec, mask):
  """Mask `vec` in log-space.

  Elements in `vec` that are not in `mask` will be set to be very
  negative, so that when no longer in log-space (e.i., after `tf.exp(vec)`)
  their values will be very close to zero.

  Args:
    vec: <float32>[...] tensor to mask
    mask: Either a None (in which case this is a no-op), a boolean or 0/1 float
      mask that matches all, or all but the last, dimensions of `vec`, or 1-D
      integer length mask

  Raises:
    ValueError: If `mask` cannot be matched to `vec`
  Returns:
    masked: vec:<float32>[...]
  """
  if mask is None:
    return vec

  if mask.dtype == tf.int32:
    # Assume `mask` holds sequence lengths
    if len(vec.shape) not in [2, 3]:
      raise ValueError("Can't use a length mask on tensor of rank>3")
    mask = tf.sequence_mask(mask, tf.shape(vec)[1], tf.float32)
  else:
    mask = tf.cast(mask, tf.float32)

  if len(mask.shape) == (len(vec.shape) - 1):
    mask = tf.expand_dims(mask, len(vec.shape) - 1)

  return vec * mask - (1 - mask) * 1E20
//  [python]
def _get_dictionary_tensor( dictionary_path, charset ):
    return tf.sparse_tensor_to_dense( tf.to_int32(
	dictionary_from_file( dictionary_path, charset )))
//  [python]
def _get_dictionary_tensor( dictionary_path, charset ):
    return tf.sparse_tensor_to_dense( tf.cast(
	dictionary_from_file( dictionary_path, charset ), tf.int32))
//  [python]
def call(self, x):
    """Get token embeddings of x.

    Args:
      x: An int64 tensor with shape [batch_size, length]
    Returns:
      embeddings: float32 tensor with shape [batch_size, length, embedding_size]
      padding: float32 tensor with shape [batch_size, length] indicating the
        locations of the padding tokens in x.
    """
    with tf.name_scope("embedding"):
      # Create binary mask of size [batch_size, length]
      mask = tf.to_float(tf.not_equal(x, 0))

      if self.method == "gather":
        embeddings = tf.gather(self.shared_weights, x)
        embeddings *= tf.expand_dims(mask, -1)
      else:  # matmul
        embeddings = tpu_utils.embedding_matmul(
            embedding_table=self.shared_weights,
            values=tf.cast(x, dtype=tf.int32),
            mask=mask
        )
        # embedding_matmul already zeros out masked positions, so
        # `embeddings *= tf.expand_dims(mask, -1)` is unnecessary.


      # Scale embedding by the sqrt of the hidden size
      embeddings *= self.hidden_size ** 0.5

      return embeddings
//  [python]
def call(self, x):
    """Get token embeddings of x.

    Args:
      x: An int64 tensor with shape [batch_size, length]
    Returns:
      embeddings: float32 tensor with shape [batch_size, length, embedding_size]
      padding: float32 tensor with shape [batch_size, length] indicating the
        locations of the padding tokens in x.
    """
    with tf.name_scope("embedding"):
      # Create binary mask of size [batch_size, length]
      mask = tf.cast(tf.not_equal(x, 0), tf.float32)

      if self.method == "gather":
        embeddings = tf.gather(self.shared_weights, x)
        embeddings *= tf.expand_dims(mask, -1)
      else:  # matmul
        embeddings = tpu_utils.embedding_matmul(
            embedding_table=self.shared_weights,
            values=tf.cast(x, dtype=tf.int32),
            mask=mask
        )
        # embedding_matmul already zeros out masked positions, so
        # `embeddings *= tf.expand_dims(mask, -1)` is unnecessary.


      # Scale embedding by the sqrt of the hidden size
      embeddings *= self.hidden_size ** 0.5

      return embeddings
//  [python]
def calc_logp_gaussian(x_tf, mean_tf, std_tf):
  dim = tf.to_float(tf.shape(x_tf)[-1])

  if mean_tf is None:
    diff_tf = x_tf
  else:
    diff_tf = x_tf - mean_tf

  logp_tf = -0.5 * tf.reduce_sum(tf.square(diff_tf / std_tf), axis=-1)
  logp_tf += -0.5 * dim * np.log(2 * np.pi) - tf.reduce_sum(tf.log(std_tf), axis=-1)

  return logp_tf
//  [python]
def calc_logp_gaussian(x_tf, mean_tf, std_tf):
  dim = tf.cast(tf.shape(x_tf)[-1], tf.float32)

  if mean_tf is None:
    diff_tf = x_tf
  else:
    diff_tf = x_tf - mean_tf

  logp_tf = -0.5 * tf.reduce_sum(tf.square(diff_tf / std_tf), axis=-1)
  logp_tf += -0.5 * dim * np.log(2 * np.pi) - tf.reduce_sum(tf.log(std_tf), axis=-1)

  return logp_tf
//  [python]
def _decode_input_tensor_to_features_dict(feature_map, hparams, decode_hp):
  """Convert the interactive input format (see above) to a dictionary.

  Args:
    feature_map: dict with inputs.
    hparams: model hyperparameters
    decode_hp: decode hyperparameters

  Returns:
    a features dictionary, as expected by the decoder.
  """
  inputs = tf.convert_to_tensor(feature_map["inputs"])
  input_is_image = False

  x = inputs
  p_hparams = hparams.problem_hparams
  # Add a third empty dimension
  x = tf.expand_dims(x, axis=[2])
  x = tf.to_int32(x)
  input_space_id = tf.constant(p_hparams.input_space_id)
  target_space_id = tf.constant(p_hparams.target_space_id)

  features = {}
  features["input_space_id"] = input_space_id
  features["target_space_id"] = target_space_id
  features["decode_length"] = (
      IMAGE_DECODE_LENGTH if input_is_image else
      tf.constant(decode_hp.extra_length))
  features["inputs"] = x
  # Save inputs to "partial_targets" when prepending inputs to targets. Also
  # keep "inputs" as some models crash if they don't exist.
  if getattr(hparams, "prepend_mode", "none") != "none":
    shape = tf.shape(x)
    partial_targets = tf.reshape(x, [shape[0], shape[1]])
    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
    features["partial_targets"] = partial_targets
  return features
//  [python]
def _decode_input_tensor_to_features_dict(feature_map, hparams, decode_hp):
  """Convert the interactive input format (see above) to a dictionary.

  Args:
    feature_map: dict with inputs.
    hparams: model hyperparameters
    decode_hp: decode hyperparameters

  Returns:
    a features dictionary, as expected by the decoder.
  """
  inputs = tf.convert_to_tensor(feature_map["inputs"])
  input_is_image = False

  x = inputs
  p_hparams = hparams.problem_hparams
  # Add a third empty dimension
  x = tf.expand_dims(x, axis=[2])
  x = tf.cast(x, tf.int32)
  input_space_id = tf.constant(p_hparams.input_space_id)
  target_space_id = tf.constant(p_hparams.target_space_id)

  features = {}
  features["input_space_id"] = input_space_id
  features["target_space_id"] = target_space_id
  features["decode_length"] = (
      IMAGE_DECODE_LENGTH if input_is_image else
      tf.constant(decode_hp.extra_length))
  features["inputs"] = x
  # Save inputs to "partial_targets" when prepending inputs to targets. Also
  # keep "inputs" as some models crash if they don't exist.
  if getattr(hparams, "prepend_mode", "none") != "none":
    shape = tf.shape(x)
    partial_targets = tf.reshape(x, [shape[0], shape[1]])
    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
    features["partial_targets"] = partial_targets
  return features
//  [python]
def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  # This needs to be kept in sync with `FeatureWriter`.
  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),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["answer_types"] = tf.FixedLenFeature([], tf.int64)

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

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

    return example

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

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

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

    return d

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

  # This needs to be kept in sync with `FeatureWriter`.
  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),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["answer_types"] = tf.FixedLenFeature([], tf.int64)

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

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

    return example

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

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

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

    return d

  return input_fn
//  [python]
def input_fn_builder(
    tfrecord_file, seq_length, is_training, batch_size, drop_remainder=True
):
    """Creates an `input_fn` closure to be passed to Estimator."""
    logging.info(
        "Creating input fun with batch_size: {} and drop remainder: {}".format(
            str(batch_size), str(drop_remainder)
        )
    )
    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "score": tf.FixedLenFeature([], tf.float32),
    }

    def _decode_record(record, name_to_features):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, name_to_features)
        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.to_int32(t)
            example[name] = t
        return example

    def input_fn(params):  # pylint: disable=unused-argument
        """Acutal data generator."""
        d = tf.data.TFRecordDataset(tfrecord_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=FLAGS.shuffle_buffer_size)
        d = d.map(lambda record: _decode_record(record, name_to_features))
        d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
        return d

    return input_fn