jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _cutout_inside_bbox(image, bbox, pad_fraction):
    """Generates cutout mask and the mean pixel value of the bbox.
    First a location is randomly chosen within the image as the center where the
    cutout mask will be applied. Note this can be towards the boundaries of the
    image, so the full cutout mask may not be applied.
    Args:
      image: 3D uint8 Tensor.
      bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
        of type float that represents the normalized coordinates between 0 and 1.
      pad_fraction: Float that specifies how large the cutout mask should be in
        in reference to the size of the original bbox. If pad_fraction is 0.25,
        then the cutout mask will be of shape
        (0.25 * bbox height, 0.25 * bbox width).
    Returns:
      A tuple. Fist element is a tensor of the same shape as image where each
      element is either a 1 or 0 that is used to determine where the image
      will have cutout applied. The second element is the mean of the pixels
      in the image where the bbox is located.
    """
    image_height = tf.shape(image)[0]
    image_width = tf.shape(image)[1]
    # Transform from shape [1, 4] to [4].
    bbox = tf.squeeze(bbox)

    min_y = tf.to_int32(tf.to_float(image_height) * bbox[0])
    min_x = tf.to_int32(tf.to_float(image_width) * bbox[1])
    max_y = tf.to_int32(tf.to_float(image_height) * bbox[2])
    max_x = tf.to_int32(tf.to_float(image_width) * bbox[3])

    # Calculate the mean pixel values in the bounding box, which will be used
    # to fill the cutout region.
    mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                          reduction_indices=[0, 1])

    # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
    # region lies entirely within the bbox.
    box_height = max_y - min_y + 1
    box_width = max_x - min_x + 1
    pad_size_height = tf.to_int32(pad_fraction * (box_height / 2))
    pad_size_width = tf.to_int32(pad_fraction * (box_width / 2))

    # Sample the center location in the image where the zero mask will be applied.
    cutout_center_height = tf.random_uniform(
        shape=[], minval=min_y, maxval=max_y + 1,
        dtype=tf.int32)

    cutout_center_width = tf.random_uniform(
        shape=[], minval=min_x, maxval=max_x + 1,
        dtype=tf.int32)

    lower_pad = tf.maximum(
        0, cutout_center_height - pad_size_height)
    upper_pad = tf.maximum(
        0, image_height - cutout_center_height - pad_size_height)
    left_pad = tf.maximum(
        0, cutout_center_width - pad_size_width)
    right_pad = tf.maximum(
        0, image_width - cutout_center_width - pad_size_width)

    cutout_shape = [image_height - (lower_pad + upper_pad),
                    image_width - (left_pad + right_pad)]
    padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

    mask = tf.pad(
        tf.zeros(cutout_shape, dtype=image.dtype),
        padding_dims, constant_values=1)

    mask = tf.expand_dims(mask, 2)
    mask = tf.tile(mask, [1, 1, 3])

    return mask, mean
//  [python]
def _cutout_inside_bbox(image, bbox, pad_fraction):
    """Generates cutout mask and the mean pixel value of the bbox.
    First a location is randomly chosen within the image as the center where the
    cutout mask will be applied. Note this can be towards the boundaries of the
    image, so the full cutout mask may not be applied.
    Args:
      image: 3D uint8 Tensor.
      bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
        of type float that represents the normalized coordinates between 0 and 1.
      pad_fraction: Float that specifies how large the cutout mask should be in
        in reference to the size of the original bbox. If pad_fraction is 0.25,
        then the cutout mask will be of shape
        (0.25 * bbox height, 0.25 * bbox width).
    Returns:
      A tuple. Fist element is a tensor of the same shape as image where each
      element is either a 1 or 0 that is used to determine where the image
      will have cutout applied. The second element is the mean of the pixels
      in the image where the bbox is located.
    """
    image_height = tf.shape(image)[0]
    image_width = tf.shape(image)[1]
    # Transform from shape [1, 4] to [4].
    bbox = tf.squeeze(bbox)

    min_y = tf.cast(tf.to_float(image_height) * bbox[0], tf.int32)
    min_x = tf.cast(tf.to_float(image_width) * bbox[1], tf.int32)
    max_y = tf.cast(tf.to_float(image_height) * bbox[2], tf.int32)
    max_x = tf.cast(tf.to_float(image_width) * bbox[3], tf.int32)

    # Calculate the mean pixel values in the bounding box, which will be used
    # to fill the cutout region.
    mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                          reduction_indices=[0, 1])

    # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
    # region lies entirely within the bbox.
    box_height = max_y - min_y + 1
    box_width = max_x - min_x + 1
    pad_size_height = tf.cast(pad_fraction * (box_height / 2), tf.int32)
    pad_size_width = tf.cast(pad_fraction * (box_width / 2), tf.int32)

    # Sample the center location in the image where the zero mask will be applied.
    cutout_center_height = tf.random_uniform(
        shape=[], minval=min_y, maxval=max_y + 1,
        dtype=tf.int32)

    cutout_center_width = tf.random_uniform(
        shape=[], minval=min_x, maxval=max_x + 1,
        dtype=tf.int32)

    lower_pad = tf.maximum(
        0, cutout_center_height - pad_size_height)
    upper_pad = tf.maximum(
        0, image_height - cutout_center_height - pad_size_height)
    left_pad = tf.maximum(
        0, cutout_center_width - pad_size_width)
    right_pad = tf.maximum(
        0, image_width - cutout_center_width - pad_size_width)

    cutout_shape = [image_height - (lower_pad + upper_pad),
                    image_width - (left_pad + right_pad)]
    padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

    mask = tf.pad(
        tf.zeros(cutout_shape, dtype=image.dtype),
        padding_dims, constant_values=1)

    mask = tf.expand_dims(mask, 2)
    mask = tf.tile(mask, [1, 1, 3])

    return mask, mean
//  [python]
def cross_entropy_seq_with_mask(logits, target_seqs, input_mask, return_details=False, name=None):
    """Returns the expression of cross-entropy of two sequences, implement
    softmax internally. Normally be used for Dynamic RNN outputs.

    Parameters
    -----------
    logits : network identity outputs
        2D tensor, ``network.outputs``, [batch_size, number of output units].
    target_seqs : int of tensor, like word ID.
        [batch_size, ?]
    input_mask : the mask to compute loss
        The same size with target_seqs, normally 0 and 1.
    return_details : boolean
        - If False (default), only returns the loss.
        - If True, returns the loss, losses, weights and targets (reshape to one vetcor).

    Examples
    --------
    - see Image Captioning Example.
    """
    targets = tf.reshape(target_seqs, [-1])  # to one vector
    weights = tf.to_float(tf.reshape(input_mask, [-1]))  # to one vector like targets
    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name) * weights
    #losses = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name)) # for TF1.0 and others

    # try: ## TF1.0
    loss = tf.divide(
        tf.reduce_sum(losses),  # loss from mask. reduce_sum before element-wise mul with mask !!
        tf.reduce_sum(weights),
        name="seq_loss_with_mask")
    # except: ## TF0.12
    #     loss = tf.div(tf.reduce_sum(losses),   # loss from mask. reduce_sum before element-wise mul with mask !!
    #                     tf.reduce_sum(weights),
    #                     name="seq_loss_with_mask")
    if return_details:
        return loss, losses, weights, targets
    else:
        return loss
//  [python]
def cross_entropy_seq_with_mask(logits, target_seqs, input_mask, return_details=False, name=None):
    """Returns the expression of cross-entropy of two sequences, implement
    softmax internally. Normally be used for Dynamic RNN outputs.

    Parameters
    -----------
    logits : network identity outputs
        2D tensor, ``network.outputs``, [batch_size, number of output units].
    target_seqs : int of tensor, like word ID.
        [batch_size, ?]
    input_mask : the mask to compute loss
        The same size with target_seqs, normally 0 and 1.
    return_details : boolean
        - If False (default), only returns the loss.
        - If True, returns the loss, losses, weights and targets (reshape to one vetcor).

    Examples
    --------
    - see Image Captioning Example.
    """
    targets = tf.reshape(target_seqs, [-1])  # to one vector
    weights = tf.cast(tf.reshape(input_mask, [-1]), tf.float32)  # to one vector like targets
    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name) * weights
    #losses = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name)) # for TF1.0 and others

    # try: ## TF1.0
    loss = tf.divide(
        tf.reduce_sum(losses),  # loss from mask. reduce_sum before element-wise mul with mask !!
        tf.reduce_sum(weights),
        name="seq_loss_with_mask")
    # except: ## TF0.12
    #     loss = tf.div(tf.reduce_sum(losses),   # loss from mask. reduce_sum before element-wise mul with mask !!
    #                     tf.reduce_sum(weights),
    #                     name="seq_loss_with_mask")
    if return_details:
        return loss, losses, weights, targets
    else:
        return loss
//  [python]
def get_splits(dataset, n_context, batch_size, points_perm=True):
  """Splits the dataset into target and context sets."""
  full_x, full_y = dataset
  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])
  context_x = target_x[:, :n_context, :]
  context_y = target_y[:, :n_context, :]
  unseen_targets = target_y[:, n_context:]
  return context_x, context_y, target_x, target_y, unseen_targets
//  [python]
def get_splits(dataset, n_context, batch_size, points_perm=True):
  """Splits the dataset into target and context sets."""
  full_x, full_y = dataset
  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)
  context_x = target_x[:, :n_context, :]
  context_y = target_y[:, :n_context, :]
  unseen_targets = target_y[:, n_context:]
  return context_x, context_y, target_x, target_y, unseen_targets
//  [python]
def _update_stats(self, x):
        ct = tf.maximum(1e-10, self._ct)

        ct_b = tf.to_float(tf.shape(x)[0])
        mu_b, var_b = tf.nn.moments(x, axes=[0])

        delta = mu_b - self._mu

        new_ct = ct + ct_b
        new_mu = self._mu + delta * ct_b / new_ct
        new_var = (self._var * ct + var_b * ct_b + delta ** 2 * ct * ct_b / new_ct) / new_ct

        self.add_update([
            tf.assign(self._ct, new_ct),
            tf.assign(self._mu, new_mu),
            tf.assign(self._var, new_var)
        ])

        return new_mu, new_var
//  [python]
def _update_stats(self, x):
        ct = tf.maximum(1e-10, self._ct)

        ct_b = tf.cast(tf.shape(x)[0], tf.float32)
        mu_b, var_b = tf.nn.moments(x, axes=[0])

        delta = mu_b - self._mu

        new_ct = ct + ct_b
        new_mu = self._mu + delta * ct_b / new_ct
        new_var = (self._var * ct + var_b * ct_b + delta ** 2 * ct * ct_b / new_ct) / new_ct

        self.add_update([
            tf.assign(self._ct, new_ct),
            tf.assign(self._mu, new_mu),
            tf.assign(self._var, new_var)
        ])

        return new_mu, new_var
//  [python]
def reconstruct(self, x, use_mode=True, use_mean=False):
    """Reconstructs the given observations.

    Args:
      x: Observed `Tensor`.
      use_mode: Boolean, if true, take the argmax over q(y|x)
      use_mean: Boolean, if true, use pixel-mean for reconstructions.

    Returns:
      The reconstructed samples x ~ p(x | y~q(y|x), z~q(z|x, y)).
    """

    hiddens = self._shared_encoder(x, is_training=self._is_training)
    qy = self.infer_cluster(hiddens)
    y_sample = qy.mode() if use_mode else qy.sample()
    y_sample = tf.to_float(y_sample)
    qz = self.infer_latent(hiddens, y_sample)
    p = self.predict(qz.sample(), y_sample)

    if use_mean:
      return p.mean()
    else:
      return p.sample()
//  [python]
def reconstruct(self, x, use_mode=True, use_mean=False):
    """Reconstructs the given observations.

    Args:
      x: Observed `Tensor`.
      use_mode: Boolean, if true, take the argmax over q(y|x)
      use_mean: Boolean, if true, use pixel-mean for reconstructions.

    Returns:
      The reconstructed samples x ~ p(x | y~q(y|x), z~q(z|x, y)).
    """

    hiddens = self._shared_encoder(x, is_training=self._is_training)
    qy = self.infer_cluster(hiddens)
    y_sample = qy.mode() if use_mode else qy.sample()
    y_sample = tf.cast(y_sample, tf.float32)
    qz = self.infer_latent(hiddens, y_sample)
    p = self.predict(qz.sample(), y_sample)

    if use_mean:
      return p.mean()
    else:
      return p.sample()
//  [python]
def sample(self, sample_shape=(), y=None, mean=False):
    """Draws a sample from the learnt distribution p(x).

    Args:
      sample_shape: `int` or 0D `Tensor` giving the number of samples to return.
        If  empty tuple (default value), 1 sample will be returned.
      y: Optional, the one hot label on which to condition the sample.
      mean: Boolean, if True the expected value of the output distribution is
        returned, otherwise samples from the output distribution.

    Returns:
      Sample tensor of shape `[B * N, ...]` where `B` is the batch size of
      the prior, `N` is the number of samples requested, and `...` represents
      the shape of the observations.

    Raises:
      ValueError: If both `sample_shape` and `n` are provided.
      ValueError: If `sample_shape` has rank > 0 or if `sample_shape`
      is an int that is < 1.
    """
    with tf.name_scope('{}_sample'.format(self.scope_name)):
      if y is None:
        y = tf.to_float(self.compute_prior().sample(sample_shape))

      if y.shape.ndims > 2:
        y = snt.MergeDims(start=0, size=y.shape.ndims - 1, name='merge_y')(y)

      z = self._latent_decoder(y, is_training=self._is_training)
      if mean:
        samples = self.predict(z.sample(), y).mean()
      else:
        samples = self.predict(z.sample(), y).sample()
    return samples
//  [python]
def sample(self, sample_shape=(), y=None, mean=False):
    """Draws a sample from the learnt distribution p(x).

    Args:
      sample_shape: `int` or 0D `Tensor` giving the number of samples to return.
        If  empty tuple (default value), 1 sample will be returned.
      y: Optional, the one hot label on which to condition the sample.
      mean: Boolean, if True the expected value of the output distribution is
        returned, otherwise samples from the output distribution.

    Returns:
      Sample tensor of shape `[B * N, ...]` where `B` is the batch size of
      the prior, `N` is the number of samples requested, and `...` represents
      the shape of the observations.

    Raises:
      ValueError: If both `sample_shape` and `n` are provided.
      ValueError: If `sample_shape` has rank > 0 or if `sample_shape`
      is an int that is < 1.
    """
    with tf.name_scope('{}_sample'.format(self.scope_name)):
      if y is None:
        y = tf.cast(self.compute_prior().sample(sample_shape), tf.float32)

      if y.shape.ndims > 2:
        y = snt.MergeDims(start=0, size=y.shape.ndims - 1, name='merge_y')(y)

      z = self._latent_decoder(y, is_training=self._is_training)
      if mean:
        samples = self.predict(z.sample(), y).mean()
      else:
        samples = self.predict(z.sample(), y).sample()
    return samples
//  [python]
def get_gamma_t(step, tail_fraction):
  """Returns an increasing momentum: step / (step + 1)."""
  s = tf.to_float(step)
  tail_fraction = tf.to_float(tail_fraction)
  c = 1. / tail_fraction

  gamma_t = c * (s - 1.) / (1. + c * (s - 1.))
  gamma_t *= (1. - tf.sqrt((1 - c) / (s * (s - 1.))) / c)

  return tf.cond(tf.equal(step, 1), lambda: 0., lambda: gamma_t)
//  [python]
def get_gamma_t(step, tail_fraction):
  """Returns an increasing momentum: step / (step + 1)."""
  s = tf.cast(step, tf.float32)
  tail_fraction = tf.cast(tail_fraction, tf.float32)
  c = 1. / tail_fraction

  gamma_t = c * (s - 1.) / (1. + c * (s - 1.))
  gamma_t *= (1. - tf.sqrt((1 - c) / (s * (s - 1.))) / c)

  return tf.cond(tf.equal(step, 1), lambda: 0., lambda: gamma_t)
//  [python]
def infer_latent(self, hiddens, y=None, use_mean_y=False):
    """Performs inference over the latent variable z.

    Args:
      hiddens: The shared encoder activations, 4D `Tensor` of size `[B, ...]`.
      y: Categorical cluster variable, `Tensor` of size `[B, ...]`.
      use_mean_y: Boolean, whether to take the mean encoding over all y.

    Returns:
      The distribution `q(z|x, y)`, which on sample produces tensors of size
      `[N, B, ...]` where `B` is the batch size of `x` and `y`, and `N` is the
      number of samples and `...` represents the shape of the latent variables.
    """
    with tf.control_dependencies([tfc.assert_rank(hiddens, 2)]):
      if y is None:
        y = tf.to_float(self.infer_cluster(hiddens).mode())

    if use_mean_y:
      # If use_mean_y, then y must be probabilities
      all_y = tf.tile(
          tf.expand_dims(tf.one_hot(tf.range(y.shape[1]), y.shape[1]), axis=1),
          multiples=[1, y.shape[0], 1])

      # Compute z KL from x (for all possible y), and keep z's around
      z_all = tf.map_fn(
          fn=lambda y: self._latent_encoder(
              hiddens, y, is_training=self._is_training).mean(),
          elems=all_y,
          dtype=tf.float32)
      return tf.einsum('ij,jik->ik', y, z_all)
    else:
      return self._latent_encoder(hiddens, y, is_training=self._is_training)
//  [python]
def infer_latent(self, hiddens, y=None, use_mean_y=False):
    """Performs inference over the latent variable z.

    Args:
      hiddens: The shared encoder activations, 4D `Tensor` of size `[B, ...]`.
      y: Categorical cluster variable, `Tensor` of size `[B, ...]`.
      use_mean_y: Boolean, whether to take the mean encoding over all y.

    Returns:
      The distribution `q(z|x, y)`, which on sample produces tensors of size
      `[N, B, ...]` where `B` is the batch size of `x` and `y`, and `N` is the
      number of samples and `...` represents the shape of the latent variables.
    """
    with tf.control_dependencies([tfc.assert_rank(hiddens, 2)]):
      if y is None:
        y = tf.cast(self.infer_cluster(hiddens).mode(), tf.float32)

    if use_mean_y:
      # If use_mean_y, then y must be probabilities
      all_y = tf.tile(
          tf.expand_dims(tf.one_hot(tf.range(y.shape[1]), y.shape[1]), axis=1),
          multiples=[1, y.shape[0], 1])

      # Compute z KL from x (for all possible y), and keep z's around
      z_all = tf.map_fn(
          fn=lambda y: self._latent_encoder(
              hiddens, y, is_training=self._is_training).mean(),
          elems=all_y,
          dtype=tf.float32)
      return tf.einsum('ij,jik->ik', y, z_all)
    else:
      return self._latent_encoder(hiddens, y, is_training=self._is_training)
//  [python]
def _random_slice(example,
                  length):
  """Extract a random slice or pad to make all sequences a fixed length.

  For example -- if one passes in [1,2,3,4] with length=2, this would return
  one of the following: [1,2], [2,3], [3,4].

  If the input is [1, 2] with length=4, this would return [1, 2, 0, 0].

  Args:
    example: Dictionary containing a single example with the "text" key. This
      "text" key should be a vector with an integer type.
    length: Length of the slice.

  Returns:
    An example containing only a fixed slice of text.
  """
  input_length = tf.shape(example["text"])[0]
  max_idx = input_length - length
  # pylint: disable=g-long-lambda
  start_idx = tf.cond(
      tf.greater(max_idx, 0), lambda: tf.random_uniform(
          [], tf.to_int32(0), tf.cast(max_idx, tf.int32), dtype=tf.int32),
      lambda: 0)
  # pylint: enable=g-long-lambda

  to_pad = tf.maximum(length - input_length, 0)
  pad_input = tf.pad(example["text"], [[0, to_pad]])
  # copy to prevent a mutation of inputs.
  example = copy.copy(example)
  example["text"] = pad_input[start_idx:start_idx + length]
  example["text"].set_shape([length])

  pad_mask = tf.pad(tf.ones([input_length]), [[0, to_pad]])
  example["mask"] = pad_mask[start_idx:start_idx + length]
  example["mask"].set_shape([length])

  return example
//  [python]
def _random_slice(example,
                  length):
  """Extract a random slice or pad to make all sequences a fixed length.

  For example -- if one passes in [1,2,3,4] with length=2, this would return
  one of the following: [1,2], [2,3], [3,4].

  If the input is [1, 2] with length=4, this would return [1, 2, 0, 0].

  Args:
    example: Dictionary containing a single example with the "text" key. This
      "text" key should be a vector with an integer type.
    length: Length of the slice.

  Returns:
    An example containing only a fixed slice of text.
  """
  input_length = tf.shape(example["text"])[0]
  max_idx = input_length - length
  # pylint: disable=g-long-lambda
  start_idx = tf.cond(
      tf.greater(max_idx, 0), lambda: tf.random_uniform(
          [], tf.cast(0, tf.int32), tf.cast(max_idx, tf.int32), dtype=tf.int32),
      lambda: 0)
  # pylint: enable=g-long-lambda

  to_pad = tf.maximum(length - input_length, 0)
  pad_input = tf.pad(example["text"], [[0, to_pad]])
  # copy to prevent a mutation of inputs.
  example = copy.copy(example)
  example["text"] = pad_input[start_idx:start_idx + length]
  example["text"].set_shape([length])

  pad_mask = tf.pad(tf.ones([input_length]), [[0, to_pad]])
  example["mask"] = pad_mask[start_idx:start_idx + length]
  example["mask"].set_shape([length])

  return example
//  [python]
def create_model(self, optimizer):
        features = tf.placeholder(tf.int32, [None, self.seq_len], name='features')
        labels = tf.placeholder(tf.int64, [None,], name='labels')

        embs = tf.Variable(self.emb_arr, dtype=tf.float32, trainable=False)
        x = tf.nn.embedding_lookup(embs, features)
        
        stacked_lstm = rnn.MultiRNNCell(
            [rnn.BasicLSTMCell(64), rnn.BasicLSTMCell(32)])
        outputs, _ = tf.nn.dynamic_rnn(stacked_lstm, x, dtype=tf.float32)
        fc1 = tf.layers.dense(inputs=outputs[:,-1,:], units=64)
        pred = tf.squeeze(tf.layers.dense(inputs=fc1, units=1))
        
        loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=labels, logits=pred)
        #optimizer = tf.train.AdamOptimizer(learning_rate=self.lr)
        grads_and_vars = optimizer.compute_gradients(loss)
        grads, _ = zip(*grads_and_vars)
        train_op = optimizer.apply_gradients(grads_and_vars, global_step=tf.train.get_global_step())
        
        correct_pred = tf.equal(tf.to_int64(tf.greater(pred,0)), labels)
        eval_metric_ops = tf.count_nonzero(correct_pred)
        
        return features, labels, train_op, grads, eval_metric_ops, loss
//  [python]
def create_model(self, optimizer):
        features = tf.placeholder(tf.int32, [None, self.seq_len], name='features')
        labels = tf.placeholder(tf.int64, [None,], name='labels')

        embs = tf.Variable(self.emb_arr, dtype=tf.float32, trainable=False)
        x = tf.nn.embedding_lookup(embs, features)
        
        stacked_lstm = rnn.MultiRNNCell(
            [rnn.BasicLSTMCell(64), rnn.BasicLSTMCell(32)])
        outputs, _ = tf.nn.dynamic_rnn(stacked_lstm, x, dtype=tf.float32)
        fc1 = tf.layers.dense(inputs=outputs[:,-1,:], units=64)
        pred = tf.squeeze(tf.layers.dense(inputs=fc1, units=1))
        
        loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=labels, logits=pred)
        #optimizer = tf.train.AdamOptimizer(learning_rate=self.lr)
        grads_and_vars = optimizer.compute_gradients(loss)
        grads, _ = zip(*grads_and_vars)
        train_op = optimizer.apply_gradients(grads_and_vars, global_step=tf.train.get_global_step())
        
        correct_pred = tf.equal(tf.cast(tf.greater(pred,0), tf.int64), labels)
        eval_metric_ops = tf.count_nonzero(correct_pred)
        
        return features, labels, train_op, grads, eval_metric_ops, loss
//  [python]
def diag_mask(n_data, pairwise_label_scores):
  """
  This is so far only used in pariwise learning-to-rank
  Args:
    n_data: a int `Tensor`.
    pairwise_label_scores: a dense `Tensor` of shape [n_data, n_data].
  Returns:
    values in pairwise_label_scores except the diagonal
    each cell contains a paiwise score difference
    only selfs/diags are 0s
  """
  mask = tf.ones([n_data, n_data]) - tf.diag(tf.ones([n_data]))
  mask = tf.cast(mask, dtype=tf.float32)
  pair_count = tf.to_float(n_data) * (tf.to_float(n_data) - 1)
  pair_count = tf.cast(pair_count, dtype=tf.float32)
  return mask, pair_count
//  [python]
def diag_mask(n_data, pairwise_label_scores):
  """
  This is so far only used in pariwise learning-to-rank
  Args:
    n_data: a int `Tensor`.
    pairwise_label_scores: a dense `Tensor` of shape [n_data, n_data].
  Returns:
    values in pairwise_label_scores except the diagonal
    each cell contains a paiwise score difference
    only selfs/diags are 0s
  """
  mask = tf.ones([n_data, n_data]) - tf.diag(tf.ones([n_data]))
  mask = tf.cast(mask, dtype=tf.float32)
  pair_count = tf.cast(n_data, dtype=tf.float32) * (tf.cast(n_data, dtype=tf.float32) - 1)
  pair_count = tf.cast(pair_count, dtype=tf.float32)
  return mask, pair_count
//  [python]
def build_trainer(self, child_model):
    """Build the train ops by connecting Controller with a Child."""
    # actor
    self.valid_loss = tf.to_float(child_model.rl_loss)
    self.valid_loss = tf.stop_gradient(self.valid_loss)
    self.valid_ppl = tf.exp(self.valid_loss)
    self.reward = REWARD_CONSTANT / self.valid_ppl

    if self.params.controller_entropy_weight:
      self.reward += self.params.controller_entropy_weight * self.sample_entropy

    # or baseline
    self.sample_log_probs = tf.reduce_sum(self.sample_log_probs)
    self.baseline = tf.Variable(0.0, dtype=tf.float32, trainable=False)
    baseline_update = tf.assign_sub(self.baseline,
                                    ((1 - self.params.controller_baseline_dec) *
                                     (self.baseline - self.reward)))

    with tf.control_dependencies([baseline_update]):
      self.reward = tf.identity(self.reward)
    self.loss = self.sample_log_probs * (self.reward - self.baseline)

    self.train_step = tf.Variable(
        0, dtype=tf.int32, trainable=False, name='train_step')
    tf_vars = [var for var in tf.trainable_variables()
               if var.name.startswith(self.name)]

    self.train_op, self.optimizer, self.grad_norm = _build_train_op(
        loss=self.loss,
        tf_vars=tf_vars,
        learning_rate=self.params.controller_learning_rate,
        train_step=self.train_step,
        num_aggregate=self.params.controller_num_aggregate)
//  [python]
def build_trainer(self, child_model):
    """Build the train ops by connecting Controller with a Child."""
    # actor
    self.valid_loss = tf.cast(child_model.rl_loss, tf.float32)
    self.valid_loss = tf.stop_gradient(self.valid_loss)
    self.valid_ppl = tf.exp(self.valid_loss)
    self.reward = REWARD_CONSTANT / self.valid_ppl

    if self.params.controller_entropy_weight:
      self.reward += self.params.controller_entropy_weight * self.sample_entropy

    # or baseline
    self.sample_log_probs = tf.reduce_sum(self.sample_log_probs)
    self.baseline = tf.Variable(0.0, dtype=tf.float32, trainable=False)
    baseline_update = tf.assign_sub(self.baseline,
                                    ((1 - self.params.controller_baseline_dec) *
                                     (self.baseline - self.reward)))

    with tf.control_dependencies([baseline_update]):
      self.reward = tf.identity(self.reward)
    self.loss = self.sample_log_probs * (self.reward - self.baseline)

    self.train_step = tf.Variable(
        0, dtype=tf.int32, trainable=False, name='train_step')
    tf_vars = [var for var in tf.trainable_variables()
               if var.name.startswith(self.name)]

    self.train_op, self.optimizer, self.grad_norm = _build_train_op(
        loss=self.loss,
        tf_vars=tf_vars,
        learning_rate=self.params.controller_learning_rate,
        train_step=self.train_step,
        num_aggregate=self.params.controller_num_aggregate)
//  [python]
def build_model(self):

		with tf.device('/gpu:0'):
			# Inputs.
			self.represenation_input, self.label_input = self.model_inputs()

			# Optional NN for latent representations.
			self.interm = self.feature_extractor(inputs=self.represenation_input, use=True)

			# Attention.
			self.weights = self.attention(inputs=self.interm, reuse=False, scope=1)

			# Classifier.
			self.prob, self.logits, self.z = self.classifier(interm=self.interm, weights=self.weights)
			self.clas = tf.to_float(tf.math.greater(self.prob, 0.5))
			
			# Losses.
			self.loss = self.loss(label=self.label_input, logits=self.logits, prob=self.prob)

			# Optimizers.
			self.trainer  = self.optimization()
//  [python]
def build_model(self):

		with tf.device('/gpu:0'):
			# Inputs.
			self.represenation_input, self.label_input = self.model_inputs()

			# Optional NN for latent representations.
			self.interm = self.feature_extractor(inputs=self.represenation_input, use=True)

			# Attention.
			self.weights = self.attention(inputs=self.interm, reuse=False, scope=1)

			# Classifier.
			self.prob, self.logits, self.z = self.classifier(interm=self.interm, weights=self.weights)
			self.clas = tf.cast(tf.math.greater(self.prob, 0.5), dtype=tf.float32)
			
			# Losses.
			self.loss = self.loss(label=self.label_input, logits=self.logits, prob=self.prob)

			# Optimizers.
			self.trainer  = self.optimization()
//  [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)
        seqs_units += self.pcoding.code(seqs_units)
        seqs_spans = tf.tile(
            tf.expand_dims(seqs_ts, 2), [1, 1, self.seqslen]) - tf.tile(
            tf.expand_dims(seqs_ts, 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=is_training)
        seqs_masks = tf.expand_dims(tf.to_float(tf.not_equal(seqs_id, 0)), 1)
        seqs_masks = tf.tile(seqs_masks, [1, tf.shape(seqs_id)[1], 1])

        # Message-Passing
        seqs_outs = seqs_units
        for i, attention in enumerate(self.list_attention):
            with tf.variable_scope("num_blocks_%d" % i):
                # Dot-Product Multi-Head Attentions
                attention_outs = attention(seqs_outs, seqs_outs, seqs_masks, seqs_spans, is_training, causality=True)
                attention_outs = tf.layers.dropout(attention_outs, rate=self.hidden_dropout_rate, training=is_training)

                # Merge Layer
                intermediate_inputs = tf.concat([attention_outs, seqs_outs], axis=-1)
                intermediate_outs = tf.layers.dense(intermediate_inputs, self.num_units, activation=tf.nn.relu)
                seqs_outs = tf.layers.dense(intermediate_outs, self.num_units)

        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)
        seqs_units += self.pcoding.code(seqs_units)
        seqs_spans = tf.tile(
            tf.expand_dims(seqs_ts, 2), [1, 1, self.seqslen]) - tf.tile(
            tf.expand_dims(seqs_ts, 1), [1, self.seqslen, 1])
        seqs_spans = tf.cast(tf.maximum(seqs_spans, 0.), tf.float32)

        # Dropout
        seqs_units = tf.layers.dropout(seqs_units, rate=self.hidden_dropout_rate, training=is_training)
        seqs_masks = tf.expand_dims(tf.cast(tf.not_equal(seqs_id, 0), tf.float32), 1)
        seqs_masks = tf.tile(seqs_masks, [1, tf.shape(seqs_id)[1], 1])

        # Message-Passing
        seqs_outs = seqs_units
        for i, attention in enumerate(self.list_attention):
            with tf.variable_scope("num_blocks_%d" % i):
                # Dot-Product Multi-Head Attentions
                attention_outs = attention(seqs_outs, seqs_outs, seqs_masks, seqs_spans, is_training, causality=True)
                attention_outs = tf.layers.dropout(attention_outs, rate=self.hidden_dropout_rate, training=is_training)

                # Merge Layer
                intermediate_inputs = tf.concat([attention_outs, seqs_outs], axis=-1)
                intermediate_outs = tf.layers.dense(intermediate_inputs, self.num_units, activation=tf.nn.relu)
                seqs_outs = tf.layers.dense(intermediate_outs, self.num_units)

        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 get_sinusoidal_embeddings(size, dim, num = 2):
    if num == 2:
        c = tf.expand_dims(tf.to_float(tf.linspace(-1.0, 1.0, size)), axis = -1)
        i = tf.to_float(tf.range(int(dim / 4)))

        peSin = tf.sin(c / (tf.pow(10000.0, 4 * i / dim)))
        peCos = tf.cos(c / (tf.pow(10000.0, 4 * i / dim)))

        peSinX = tf.tile(tf.expand_dims(peSin, axis = 0), [size, 1, 1])
        peCosX = tf.tile(tf.expand_dims(peCos, axis = 0), [size, 1, 1])
        peSinY = tf.tile(tf.expand_dims(peSin, axis = 1), [1, size, 1])
        peCosY = tf.tile(tf.expand_dims(peCos, axis = 1), [1, size, 1])

        emb = tf.concat([peSinX, peCosX, peSinY, peCosY], axis = -1)
    else:
        pi = tf.constant(math.pi)
        theta = tf.range(0, pi, pi / num)
        dirs = tf.stack([tf.cos(theta), tf.sin(theta)], axis = -1)

        c = tf.linspace(-1.0, 1.0, size)
        x = tf.tile(tf.expand_dims(c, axis = 0), [size, 1])
        y = tf.tile(tf.expand_dims(c, axis = 1), [1, size])
        xy = tf.stack([x,y], axis = -1)

        lens = tf.reduce_sum(tf.expand_dims(xy, axis = -2) * dirs, axis = -1, keepdims = True)

        i = tf.to_float(tf.range(int(dim / (2 * num))))
        sins = tf.sin(lens / (tf.pow(10000.0, 2 * num * i / dim)))
        coss = tf.cos(lens / (tf.pow(10000.0, 2 * num * i / dim)))
        emb = tf.reshape(tf.concat([sins, coss], axis = -1), [size, size, dim])

    return emb
//  [python]
def get_sinusoidal_embeddings(size, dim, num = 2):
    if num == 2:
        c = tf.expand_dims(tf.cast(tf.linspace(-1.0, 1.0, size), dtype=tf.float32), axis = -1)
        i = tf.cast(tf.range(int(dim / 4)), dtype=tf.float32)

        peSin = tf.sin(c / (tf.pow(10000.0, 4 * i / dim)))
        peCos = tf.cos(c / (tf.pow(10000.0, 4 * i / dim)))

        peSinX = tf.tile(tf.expand_dims(peSin, axis = 0), [size, 1, 1])
        peCosX = tf.tile(tf.expand_dims(peCos, axis = 0), [size, 1, 1])
        peSinY = tf.tile(tf.expand_dims(peSin, axis = 1), [1, size, 1])
        peCosY = tf.tile(tf.expand_dims(peCos, axis = 1), [1, size, 1])

        emb = tf.concat([peSinX, peCosX, peSinY, peCosY], axis = -1)
    else:
        pi = tf.constant(math.pi)
        theta = tf.range(0, pi, pi / num)
        dirs = tf.stack([tf.cos(theta), tf.sin(theta)], axis = -1)

        c = tf.linspace(-1.0, 1.0, size)
        x = tf.tile(tf.expand_dims(c, axis = 0), [size, 1])
        y = tf.tile(tf.expand_dims(c, axis = 1), [1, size])
        xy = tf.stack([x,y], axis = -1)

        lens = tf.reduce_sum(tf.expand_dims(xy, axis = -2) * dirs, axis = -1, keepdims = True)

        i = tf.cast(tf.range(int(dim / (2 * num))), dtype=tf.float32)
        sins = tf.sin(lens / (tf.pow(10000.0, 2 * num * i / dim)))
        coss = tf.cos(lens / (tf.pow(10000.0, 2 * num * i / dim)))
        emb = tf.reshape(tf.concat([sins, coss], axis = -1), [size, size, dim])

    return emb
//  [python]
def learning_rate_schedule(base_learning_rate, num_examples):
  """Build learning rate schedule."""
  global_step = tf.train.get_or_create_global_step()
  warmup_steps = int(round(
      FLAGS.warmup_epochs * num_examples // FLAGS.train_batch_size))
  scaled_lr = base_learning_rate * FLAGS.train_batch_size / 256.0
  learning_rate = (tf.to_float(global_step) / int(warmup_steps) * scaled_lr
                   if warmup_steps else scaled_lr)

  # Cosine decay learning rate schedule
  total_steps = get_train_steps(num_examples)
  learning_rate = tf.where(
      global_step < warmup_steps, learning_rate,
      tf.train.cosine_decay(
          scaled_lr,
          global_step - warmup_steps,
          total_steps - warmup_steps))

  return learning_rate
//  [python]
def learning_rate_schedule(base_learning_rate, num_examples):
  """Build learning rate schedule."""
  global_step = tf.train.get_or_create_global_step()
  warmup_steps = int(round(
      FLAGS.warmup_epochs * num_examples // FLAGS.train_batch_size))
  scaled_lr = base_learning_rate * FLAGS.train_batch_size / 256.0
  learning_rate = (tf.cast(global_step, tf.float32) / int(warmup_steps) * scaled_lr
                   if warmup_steps else scaled_lr)

  # Cosine decay learning rate schedule
  total_steps = get_train_steps(num_examples)
  learning_rate = tf.where(
      global_step < warmup_steps, learning_rate,
      tf.train.cosine_decay(
          scaled_lr,
          global_step - warmup_steps,
          total_steps - warmup_steps))

  return learning_rate
//  [python]
def word_to_char_ids(word, word_length):
  """Convert a string to a padded vector of character ids.

  If the true length of the word is less than `word_length`, padding is added.
  If the true length of the word is greater than `word_length`, additional bytes
  are ignored.

  Args:
    word: <string> []
    word_length: Number of bytes to include per word.

  Returns:
    char_ids: <int32> [word_length]
  """
  char_ids = tf.to_int32(tf.decode_raw(word, tf.uint8))[:word_length - 2]
  padding = tf.fill([word_length - tf.shape(char_ids)[0] - 2], PAD_CHAR)
  char_ids = tf.concat([[BOW_CHAR], char_ids, [EOW_CHAR], padding], 0)
  char_ids.set_shape([word_length])
  return char_ids
//  [python]
def word_to_char_ids(word, word_length):
  """Convert a string to a padded vector of character ids.

  If the true length of the word is less than `word_length`, padding is added.
  If the true length of the word is greater than `word_length`, additional bytes
  are ignored.

  Args:
    word: <string> []
    word_length: Number of bytes to include per word.

  Returns:
    char_ids: <int32> [word_length]
  """
  char_ids = tf.cast(tf.decode_raw(word, tf.uint8), tf.int32)[:word_length - 2]
  padding = tf.fill([word_length - tf.shape(char_ids)[0] - 2], PAD_CHAR)
  char_ids = tf.concat([[BOW_CHAR], char_ids, [EOW_CHAR], padding], 0)
  char_ids.set_shape([word_length])
  return char_ids
//  [python]
def blend(image1, image2, factor):
  """Blends image1 and image2 using 'factor'.

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

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

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

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

  difference = image2 - image1
  scaled = factor * difference

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

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

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

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

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

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

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

  difference = image2 - image1
  scaled = factor * difference

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

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

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scales the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.to_float(tf.reduce_min(image))
    hi = tf.to_float(tf.reduce_max(image))

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.to_float(im) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scales the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.cast(tf.reduce_min(image), tf.float32)
    hi = tf.cast(tf.reduce_max(image), tf.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.cast(im, tf.float32) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

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

    if is_training:
        name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
        name_to_features["end_positions"] = 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