jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def prepare_pegasus_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = t.cast(t.math.not_equal(input_ids, config.pad_token_id), t.int8)
    if decoder_attention_mask is None:
        decoder_attention_mask = t.concat(
            [
                t.ones(decoder_input_ids[:, :1].shape, dtype=t.int8),
                t.cast(t.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), t.int8),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = t.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = t.ones((config.decoder_layers, config.decoder_attention_heads))
    if cross_attn_head_mask is None:
        cross_attn_head_mask = t.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }
//  [python]
def prepare_pegasus_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8)
    if decoder_attention_mask is None:
        decoder_attention_mask = tf.concat(
            [
                tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8),
                tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = tf.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    if cross_attn_head_mask is None:
        cross_attn_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFPegasusModel(config=config).get_decoder()
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        head_mask = inputs_dict["head_mask"]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, head_mask=head_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = tf.to_int32(ids_tensor((self.batch_size, 3), 2))

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFPegasusModel(config=config).get_decoder()
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        head_mask = inputs_dict["head_mask"]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, head_mask=head_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8)

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def prepare_gt(self, batch_size, mixup=False):
    if mixup:
      gt = np.concatenate([np.eye(batch_size), np.eye(batch_size)], axis=0)
      gt = np.concatenate([gt, gt], axis=1)
      gt = np.flipud(gt) / 2 + gt
      gt = tf.convert_to_tensor(gt, dtype=tf.bool)
      return gt

    gt = tf.concat([tf.eye(batch_size), tf.eye(batch_size)], axis=0)
    gt = tf.concat([gt, gt], axis=1)
    return tf.to_float(gt)
//  [python]
def prepare_gt(self, batch_size, mixup=False):
    if mixup:
      gt = np.concatenate([np.eye(batch_size), np.eye(batch_size)], axis=0)
      gt = np.concatenate([gt, gt], axis=1)
      gt = np.flipud(gt) / 2 + gt
      gt = tf.convert_to_tensor(gt, dtype=tf.bool)
      return gt

    gt = tf.concat([tf.eye(batch_size), tf.eye(batch_size)], axis=0)
    gt = tf.concat([gt, gt], axis=1)
    return tf.cast(gt, dtype=tf.bool)
//  [python]
def sample_binomial(p,n,eps=S('binom.log_eps')):
    # sample from a binomial distribution
    if S("binom.probability_mode") == "tfp":
        P = tf.stack([p,1.0-p],axis=-1)
        weight_binom = tfp.distributions.Multinomial(total_count=n,probs=P).sample()[...,0]
        # weight_binom = tfp.distributions.Binomial(total_count=n,probs=p).sample()
        weight_binom = tf.cast(weight_binom,tf.float32)
    elif S("binom.probability_mode") == "gumbel":
        with tf.variable_scope("p"):
            # p = weight_p
            p = tf.clip_by_value(p, 0.0, 1.0)
            P = tf.stack([binomialCoeff(n,k)*p**k*(1-p)**(n-k) for k in range(n+1)],axis=-1)

            # reduces numerical instabilities
            P = tf.clip_by_value(P,eps,1.0)
            gumbel = -tf.log(tf.maximum(-tf.log(tf.maximum(tf.random.uniform(P.get_shape()),eps)),eps))

            # gumbel = -tf.log(-tf.log(tf.random.uniform(P.get_shape())))
            # tf.summary.histogram("binom_p",p)
            # tf.summary.histogram("binom_P",P)
            # tf.summary.histogram("binom_logP",tf.log(P))
        weight_binom = tf.argmax(tf.log(P)+gumbel,axis=-1)
        weight_binom = tf.cast(weight_binom,tf.float32)
    elif S("binom.probability_mode") == "gumbel_log":
        with tf.variable_scope("p"):
            # p = weight_p
            p = tf.clip_by_value(p, eps, 1.0-eps)
            logP = tf.stack([np.log(binomialCoeff(n,k))+k*tf.log(p)+(n-k)*tf.log(1-p) for k in range(n+1)],axis=-1)

            # reduces numerical instabilities
            gumbel = -tf.log(tf.maximum(-tf.log(tf.maximum(tf.random.uniform(logP.get_shape()),eps)),eps))

        weight_binom = tf.argmax(logP+gumbel,axis=-1)
        weight_binom = tf.cast(weight_binom,tf.float32)

    if S("binom.gradient_correction") == "pass":
        weight_binom = pass_gradient(p, lambda p: weight_binom, lambda p: n*p)
    elif S("binom.gradient_correction") == "gumbel":
        weight_binom = pass_gradient(p, lambda p: weight_binom, lambda p: tf.squeeze(tf.batch_gather(P,tf.cast(tf.expand_dims(weight_binom,-1),tf.int32))))
    else:
        raise ValueError("Gradient not defined for tf.cast. TODO")
    return weight_binom
//  [python]
def test_within_image_supervised_pixel_contrastive_loss(self):
    loss = contrastive_loss.within_image_supervised_pixel_contrastive_loss(
        features=tf.to_float(self.labels),
        labels=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)
    expected_loss = 2.4503216
    self.assertAllClose(loss, expected_loss)
//  [python]
def test_within_image_supervised_pixel_contrastive_loss(self):
    loss = contrastive_loss.within_image_supervised_pixel_contrastive_loss(
        features=tf.cast(self.labels, tf.float32),
        labels=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)
    expected_loss = 2.4503216
    self.assertAllClose(loss, expected_loss)
//  [python]
def test_contrastive_loss(self):
    logits = t.cast(
        t.matmul(self.labels, self.labels, transpose_b=True), dtype=t.float32)

    positive_mask = t.cast(
        t.equal(self.labels, t.transpose(self.labels, [0, 2, 1])), t.float32)
    negative_mask = 1 - positive_mask

    ignore_mask = t.math.reduce_any(
        t.equal(self.labels,
                 t.constant(self.ignore_labels, dtype=self.labels.dtype)),
        axis=2, keepdims=True)
    ignore_mask = t.cast(
        t.logical_or(ignore_mask, t.transpose(ignore_mask, [0, 2, 1])),
        t.float32)

    loss = contrastive_loss.compute_contrastive_loss(
        logits, positive_mask, negative_mask, ignore_mask)
    expected_loss = 2.4503216
    self.assertAllClose(loss, expected_loss)
//  [python]
def test_contrastive_loss(self):
    logits = tf.cast(
        tf.matmul(self.labels, self.labels, transpose_b=True), dtype=tf.float32)

    positive_mask = tf.cast(
        tf.equal(self.labels, tf.transpose(self.labels, [0, 2, 1])), tf.float32)
    negative_mask = 1 - positive_mask

    ignore_mask = tf.math.reduce_any(
        tf.equal(self.labels,
                 tf.constant(self.ignore_labels, dtype=self.labels.dtype)),
        axis=2, keepdims=True)
    ignore_mask = tf.cast(
        tf.logical_or(ignore_mask, tf.transpose(ignore_mask, [0, 2, 1])),
        tf.float32)

    loss = contrastive_loss.compute_contrastive_loss(
        logits, positive_mask, negative_mask, ignore_mask)
    expected_loss = 2.4503216
    self.assertAllClose(loss, expected_loss)
//  [python]
def ragged_top_k(
    num_samples: tf.Tensor,
    weights: tf.RaggedTensor,
    *,
    global_indices: bool,
) -> tf.RaggedTensor:
  """Implements `ext_ops.py:ragged_top_k()`."""

  def fn(values: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
    num_samples, weights = values
    res = tf.math.top_k(
        weights, k=t.cast(num_samples, dtype=tf.int32)
    )
    # TODO: b/294224429 - Remove this cast and use index_type when we stop
    # supporting tf < 2.13
    return t.cast(res.indices, dtype=num_samples.dtype)

  indices_dtype = num_samples.dtype
  row_splits = t.cast(weights.row_splits, dtype=indices_dtype)
  row_starts, row_limits = row_splits[:-1], row_splits[1:]
  row_lengths = row_limits - row_starts
  num_samples = tf.math.minimum(num_samples, row_lengths)
  result = tf.map_fn(
      fn,
      (num_samples, weights),
      fn_output_signature=tf.RaggedTensorSpec(
          [None],
          dtype=indices_dtype,
          ragged_rank=0,
          row_splits_dtype=weights.row_splits.dtype,
      ),
  )
  if global_indices:
    result += tf.expand_dims(row_starts, axis=-1)
  return result
//  [python]
def ragged_top_k(
    num_samples: tf.Tensor,
    weights: tf.RaggedTensor,
    *,
    global_indices: bool,
) -> tf.RaggedTensor:
  """Implements `ext_ops.py:ragged_top_k()`."""

  def fn(values: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
    num_samples, weights = values
    res = tf.math.top_k(
        weights, k=tf.cast(num_samples, dtype=tf.int32)
    )
    # TODO: b/294224429 - Remove this cast and use index_type when we stop
    # supporting tf < 2.13
    return tf.cast(res.indices, dtype=num_samples.dtype)

  indices_dtype = num_samples.dtype
  row_splits = tf.cast(weights.row_splits, dtype=indices_dtype)
  row_starts, row_limits = row_splits[:-1], row_splits[1:]
  row_lengths = row_limits - row_starts
  num_samples = tf.math.minimum(num_samples, row_lengths)
  result = tf.map_fn(
      fn,
      (num_samples, weights),
      fn_output_signature=tf.RaggedTensorSpec(
          [None],
          dtype=indices_dtype,
          ragged_rank=0,
          row_splits_dtype=weights.row_splits.dtype,
      ),
  )
  if global_indices:
    result += tf.expand_dims(row_starts, axis=-1)
  return result
//  [python]
def test_cross_image_supervised_pixel_contrastive_loss(self):
    loss = contrastive_loss.cross_image_supervised_pixel_contrastive_loss(
        features1=tf.to_float(self.labels),
        features2=tf.to_float(2 * self.labels),
        labels1=self.labels,
        labels2=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)
    expected_loss = 5.46027374
    self.assertAllClose(loss, expected_loss)
//  [python]
def test_cross_image_supervised_pixel_contrastive_loss(self):
    loss = contrastive_loss.cross_image_supervised_pixel_contrastive_loss(
        features1=tf.cast(self.labels, tf.float32),
        features2=tf.cast(2 * self.labels, tf.float32),
        labels1=self.labels,
        labels2=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)
    expected_loss = 5.46027374
    self.assertAllClose(loss, expected_loss)
//  [python]
def test_symmetry_of_cross_image_supervised_pixel_contrastive_loss(self):
    features1 = t.cast(self.labels, tf.float32)
    features2 = t.cast(2 * self.labels, tf.float32)

    loss1 = contrastive_loss.cross_image_supervised_pixel_contrastive_loss(
        features1=features1,
        features2=features2,
        labels1=self.labels,
        labels2=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)

    loss2 = contrastive_loss.cross_image_supervised_pixel_contrastive_loss(
        features1=features2,
        features2=features1,
        labels1=self.labels,
        labels2=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)
    self.assertAllClose(loss1, loss2)
//  [python]
def test_symmetry_of_cross_image_supervised_pixel_contrastive_loss(self):
    features1 = tf.cast(self.labels, tf.float32)
    features2 = tf.cast(2 * self.labels, tf.float32)

    loss1 = contrastive_loss.cross_image_supervised_pixel_contrastive_loss(
        features1=features1,
        features2=features2,
        labels1=self.labels,
        labels2=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)

    loss2 = contrastive_loss.cross_image_supervised_pixel_contrastive_loss(
        features1=features2,
        features2=features1,
        labels1=self.labels,
        labels2=self.labels,
        ignore_labels=self.ignore_labels,
        temperature=1.0)
    self.assertAllClose(loss1, loss2)
//  [python]
def ssl_loss(embeddings,
             labels,
             infonce_temperature=0.05,
             entropy_weight=30,
             mixup=False):
  """Calculate a contrastive loss added by entropy."""

  # Calculate cosine similairty,
  similarity_matrix = embeddings @ tf.transpose(embeddings, perm=[1, 0])

  # Get logits.
  logits = tf.exp(similarity_matrix / infonce_temperature)

  # Get Match matrix based on the provided labels.
  match_matrix = tf.equal(
      tf.identity(labels),
      tf.constant(True, dtype=tf.bool))  # Copied items are set to True in the same line (Batch).
  non_matches_matrix = tf.logical_not(match_matrix)
  non_matches_matrix = tf.cast(non_matches_matrix, tf.float32)
  identity = tf.eye(labels.shape[0], dtype=tf.bool)
  nontrivial_matches = tf.cast(
      tf.logical_and(labels, tf.logical_not(identity)), tf.float32)

  # Calculate partitions.
  partitions = logits + (tf.reduce_sum(non_matches_matrix * logits, axis=1) +
                         eps)  # Sum of all logits per line
  probabilities = logits / partitions  # This is not summing one???

  # Calculate InfoNCE loss.
  batch_size = similarity_matrix.shape[0]
  if mixup:

    nontrivial_sum = tf.reduce_sum(tf.cast(nontrivial_matches, tf.float32))
    infonce_loss = tf.reduce_sum(
        -tf.math.log(probabilities) * nontrivial_matches) / nontrivial_sum
  else:
    infonce_loss = tf.reduce_sum(
        -tf.math.log(probabilities) * nontrivial_matches) / batch_size

  # Calculate entropy_loss.
  small_value = tf.constant(
      -100.0, dtype=tf.float32)  # any value > max L2 normalized distance.
  max_non_match_sim = tf.reduce_max(
      tf.where(
          tf.cast(non_matches_matrix, tf.bool), similarity_matrix, small_value),
      axis=1  #, keepdims=True
  )
  closest_distance = tf.sqrt(
      tf.clip_by_value(
          2 - (2 * max_non_match_sim), clip_value_min=1e-6,
          clip_value_max=1)  # Max = 1 since it is a similairity score.
  )
  entropy_loss = -tf.reduce_mean(
      tf.math.log(closest_distance), axis=0) * entropy_weight

  loss = infonce_loss + entropy_loss
  return loss, infonce_loss, entropy_loss
//  [python]
def ssl_loss(embeddings,
             labels,
             infonce_temperature=0.05,
             entropy_weight=30,
             mixup=False):
  """Calculate a contrastive loss added by entropy."""

  # Calculate cosine similairty,
  similarity_matrix = embeddings @ tf.transpose(embeddings, perm=[1, 0])

  # Get logits.
  logits = tf.exp(similarity_matrix / infonce_temperature)

  # Get Match matrix based on the provided labels.
  match_matrix = tf.cast(
      tf.identity(labels),
      dtype=tf.bool)  # Copied items are set to True in the same line (Batch).
  non_matches_matrix = ~match_matrix
  non_matches_matrix = tf.cast(non_matches_matrix, tf.float32)
  identity = tf.eye(labels.shape[0], dtype=tf.bool)
  nontrivial_matches = tf.cast(
      tf.math.logical_and(labels, (~identity)), tf.float32)

  # Calculate partitions.
  partitions = logits + (tf.reduce_sum(non_matches_matrix * logits, axis=1) +
                         eps)  # Sum of all logits per line
  probabilities = logits / partitions  # This is not summing one???

  # Calculate InfoNCE loss.
  batch_size = similarity_matrix.shape[0]
  if mixup:

    nontrivial_sum = tf.reduce_sum(tf.cast(nontrivial_matches, tf.float32))
    infonce_loss = tf.reduce_sum(
        -tf.math.log(probabilities) * nontrivial_matches) / nontrivial_sum
  else:
    infonce_loss = tf.reduce_sum(
        -tf.math.log(probabilities) * nontrivial_matches) / batch_size

  # Calculate entropy_loss.
  small_value = tf.constant(
      -100.0, dtype=tf.float32)  # any value > max L2 normalized distance.
  max_non_match_sim = tf.reduce_max(
      tf.where(
          tf.cast(non_matches_matrix, tf.bool), similarity_matrix, small_value),
      axis=1  #, keepdims=True
  )
  closest_distance = tf.sqrt(
      tf.clip_by_value(
          2 - (2 * max_non_match_sim), clip_value_min=1e-6,
          clip_value_max=1)  # Max = 1 since it is a similairity score.
  )
  entropy_loss = -tf.reduce_mean(
      tf.math.log(closest_distance), axis=0) * entropy_weight

  loss = infonce_loss + entropy_loss
  return loss, infonce_loss, entropy_loss
//  [python]
def ssl_local_loss(embs, t_embs, t_indices, temperature=0.5):
  """Calculate the local loss.

  Args:
    embs:
    t_embs:
    t_indices:
    temperature:

  Returns:

  """
  # Reshape input.
  # pylint: disable=invalid-name
  B, _, _, D = embs.shape
  # pylint: enable=invalid-name
  embs = tf.reshape(embs, [B, -1, D])
  t_embs = tf.reshape(t_embs, [B, -1, D])
  t_indices = tf.reshape(t_indices, [B, -1])

  # Normalize embs.
  # pylint: disable=invalid-name
  L2Norm = L2Normalization()
  # pylint: enable=invalid-name
  embs = L2Norm(embs, axis=2)
  t_embs = L2Norm(t_embs, axis=2)
  t_indices = tf.cast(t_indices, tf.int32)

  valid_inds = t_indices >= 0

  # Fix indices.
  # Invalid index were set to -1.
  t_indices = tf.where(valid_inds, t_indices, 0)
  maped_embs = tf.gather(embs, t_indices, batch_dims=1)

  # Calculate Similarity matrix.
  sim_matrix = maped_embs @ tf.transpose(
      t_embs, perm=[0, 2, 1])  # B, nBlock*nBlocks

  logits = tf.exp(sim_matrix / temperature)

  matches = tf.eye(
      num_rows=sim_matrix.shape[1],
      num_columns=sim_matrix.shape[2],
      batch_shape=[sim_matrix.shape[0]],
      dtype=tf.bool)

  not_matches = tf.cast(~matches, tf.float32)
  # Only consider valid tranformed pixels.
  # i.e. pixels that actually were transformted from imgA to imgB.
  # Valid_inds=(B, M*N).
  valid_inds = tf.expand_dims(valid_inds, axis=2)

  logits = tf.where(valid_inds, logits, 0)

  # Calculate partitions.
  partitions = tf.expand_dims(
      tf.reduce_sum(not_matches * logits, axis=1) + eps, axis=1)
  partitions = tf.where(valid_inds, partitions, eps)

  # Validate Blocks that should be matched with the augmented masks.
  matches = tf.where(valid_inds, not_matches == 0, False)
  matches = tf.cast(matches, tf.float32)

  # Trying to Smooth likehood of each block with gaussian filter.
  # matches = gaussian_filter2d(
  #        tf.reshape(matches, (B*M*N, M, N, 1)),
  #        filter_shape=3,
  #        sigma=0.425)
  # matches = tf.reshape(matches, (B, M*N, M*N))

  # Calculate probs.
  probabilities = logits / (partitions + eps)
  matches_probs = -tf.math.log(probabilities) * matches
  # Considerate only valid probs
  matches_probs = tf.where(valid_inds, matches_probs, 0)

  valid_inds = tf.cast(valid_inds, tf.float32)

  loss = tf.reduce_sum(matches_probs) / (tf.reduce_sum(valid_inds) + eps)
  return loss
//  [python]
def test_tf_mix_jax_with_uncompilableble(self):
    """Show how to combine TF-uncompilableble code with compiled JAX-converted code."""
    def tf_fn(x_str, compute_tf_fn=lambda x: x):
      # Some TF preprocessing code that cannot be compiled with XLA because it
      # uses strings.
      numbers_f32 = tf.strings.to_number(x_str, out_type=tf.float32)
      numbers_f16 = t.cast(numbers_f32, tf.float16)
      return compute_tf_fn(numbers_f16)

    x_str = np.array(["3.14", "2.78"])

    # Test that we get an error if we try to TF-compile `tf_fn`
    with self.assertRaisesRegex(
        Exception,
        "Detected unsupported operations when trying to compile graph"):
      tf.function(tf_fn, jit_compile=True, autograph=False)(x_str)

    # Plug in the TF-compiled JAX-converted `compute_jax_fn`.
    composed_fn = lambda x_str: tf_fn(
        x_str,
        compute_tf_fn=tf.function(jax2tf.convert(jnp.sin),
                                  autograph=False,
                                  jit_compile=True))
    res_tf = composed_fn(x_str)
    self.assertAllClose(res_tf.numpy(),
                        jnp.sin(np.array([3.14, 2.78], dtype=np.float16)))

    # Save and restore SavedModel
    restored_f, _ = tf_test_util.SaveAndLoadFunction(composed_fn,
                                                     input_args=[x_str])
    res_tf_restored = restored_f(x_str)
    self.assertAllClose(res_tf_restored.numpy(), res_tf.numpy())
//  [python]
def test_tf_mix_jax_with_uncompilableble(self):
    """Show how to combine TF-uncompilableble code with compiled JAX-converted code."""
    def tf_fn(x_str, compute_tf_fn=lambda x: x):
      # Some TF preprocessing code that cannot be compiled with XLA because it
      # uses strings.
      numbers_f32 = tf.strings.to_number(x_str, out_type=tf.float32)
      numbers_f16 = tf.cast(numbers_f32, tf.float16)
      return compute_tf_fn(numbers_f16)

    x_str = np.array(["3.14", "2.78"])

    # Test that we get an error if we try to TF-compile `tf_fn`
    with self.assertRaisesRegex(
        Exception,
        "Detected unsupported operations when trying to compile graph"):
      tf.function(tf_fn, jit_compile=True, autograph=False)(x_str)

    # Plug in the TF-compiled JAX-converted `compute_jax_fn`.
    composed_fn = lambda x_str: tf_fn(
        x_str,
        compute_tf_fn=tf.function(jax2tf.convert(jnp.sin),
                                  autograph=False,
                                  jit_compile=True))
    res_tf = composed_fn(x_str)
    self.assertAllClose(res_tf.numpy(),
                        jnp.sin(np.array([3.14, 2.78], dtype=np.float16)))

    # Save and restore SavedModel
    restored_f, _ = tf_test_util.SaveAndLoadFunction(composed_fn,
                                                     input_args=[x_str])
    res_tf_restored = restored_f(x_str)
    self.assertAllClose(res_tf_restored.numpy(), res_tf.numpy())
//  [python]
def normalize_image(
    frames: tf.Tensor,
    zero_centering_image: bool,
    dtype: tf.dtypes.DType = tf.float32,
) -> tf.Tensor:
  """Normalizes images.

  Args:
    frames: A Tensor of numbers.
    zero_centering_image: If True, results are in [-1, 1], if False, results are
      in [0, 1].
    dtype: Type of output Tensor.

  Returns:
    A Tensor of same shape as the input and of the given type.
  """
  frames = t.cast(frames, dtype)
  if zero_centering_image:
    return frames * (2.0 / 255.0) - 1.0
  else:
    return frames / 255.0
//  [python]
def normalize_image(
    frames: tf.Tensor,
    zero_centering_image: bool,
    dtype: tf.dtypes.DType = tf.float32,
) -> tf.Tensor:
  """Normalizes images.

  Args:
    frames: A Tensor of numbers.
    zero_centering_image: If True, results are in [-1, 1], if False, results are
      in [0, 1].
    dtype: Type of output Tensor.

  Returns:
    A Tensor of same shape as the input and of the given type.
  """
  frames = tf.cast(frames, dtype)
  if zero_centering_image:
    return frames * (2.0 / 255.0) - 1.0
  else:
    return frames / 255.0
//  [python]
def sample_linspace_sequence(sequence: tf.Tensor, num_windows: int,
                             num_steps: int, stride: int) -> tf.Tensor:
  """Samples `num_windows` segments from sequence with linearly spaced offsets.

  The samples are concatenated in a single `tf.Tensor` in order to have the same
  format structure per timestep (e.g. a single frame). If `num_steps` * `stride`
  is bigger than the number of timesteps, the sequence is repeated. This
  function can be used in evaluation in order to extract enough segments to span
  the entire sequence.

  Args:
    sequence: Any tensor where the first dimension is timesteps.
    num_windows: Number of windows retrieved from the sequence.
    num_steps: Number of steps (e.g. frames) to take.
    stride: Distance to sample between timesteps.

  Returns:
    A single `tf.Tensor` with first dimension `num_windows` * `num_steps`. The
    tensor contains the concatenated list of `num_windows` tensors which offsets
    have been linearly spaced from input.
  """
  sequence_length = tf.shape(sequence)[0]
  max_offset = tf.maximum(0, sequence_length - num_steps * stride)
  offsets = tf.linspace(0.0, tf.cast(max_offset, tf.float32), num_windows)
  offsets = tf.to_int32(offsets)

  all_indices = []
  for i in range(num_windows):
    all_indices.append(
        _sample_or_pad_sequence_indices(
            sequence=sequence,
            num_steps=num_steps,
            stride=stride,
            offset=offsets[i]))

  indices = tf.concat(all_indices, axis=0)
  indices.set_shape((num_windows * num_steps,))
  return tf.gather(sequence, indices)
//  [python]
def sample_linspace_sequence(sequence: tf.Tensor, num_windows: int,
                             num_steps: int, stride: int) -> tf.Tensor:
  """Samples `num_windows` segments from sequence with linearly spaced offsets.

  The samples are concatenated in a single `tf.Tensor` in order to have the same
  format structure per timestep (e.g. a single frame). If `num_steps` * `stride`
  is bigger than the number of timesteps, the sequence is repeated. This
  function can be used in evaluation in order to extract enough segments to span
  the entire sequence.

  Args:
    sequence: Any tensor where the first dimension is timesteps.
    num_windows: Number of windows retrieved from the sequence.
    num_steps: Number of steps (e.g. frames) to take.
    stride: Distance to sample between timesteps.

  Returns:
    A single `tf.Tensor` with first dimension `num_windows` * `num_steps`. The
    tensor contains the concatenated list of `num_windows` tensors which offsets
    have been linearly spaced from input.
  """
  sequence_length = tf.shape(sequence)[0]
  max_offset = tf.maximum(0, sequence_length - num_steps * stride)
  offsets = tf.linspace(0.0, tf.cast(max_offset, tf.float32), num_windows)
  offsets = tf.cast(offsets, tf.int32)

  all_indices = []
  for i in range(num_windows):
    all_indices.append(
        _sample_or_pad_sequence_indices(
            sequence=sequence,
            num_steps=num_steps,
            stride=stride,
            offset=offsets[i]))

  indices = tf.concat(all_indices, axis=0)
  indices.set_shape((num_windows * num_steps,))
  return tf.gather(sequence, indices)
//  [python]
def resize_smallest(frames: tf.Tensor, min_resize: int) -> tf.Tensor:
  """Resizes frames so that min(`height`, `width`) is equal to `min_resize`.

  This function will not do anything if the min(`height`, `width`) is already
  equal to `min_resize`. This allows to save compute time.

  Args:
    frames: A Tensor of dimension [timesteps, input_h, input_w, channels].
    min_resize: Minimum size of the final image dimensions.

  Returns:
    A Tensor of shape [timesteps, output_h, output_w, channels] of type
      frames.dtype where min(output_h, output_w) = min_resize.
  """
  shape = tf.shape(frames)
  input_h = shape[1]
  input_w = shape[2]

  output_h = tf.maximum(min_resize, (input_h * min_resize) // input_w)
  output_w = tf.maximum(min_resize, (input_w * min_resize) // input_h)

  def resize_fn():
    frames_resized = tf.image.resize(frames, (output_h, output_w))
    return tf.to_float(frames_resized)

  should_resize = tf.math.logical_or(
      tf.not_equal(input_w, output_w), tf.not_equal(input_h, output_h))
  frames = tf.cond(should_resize, resize_fn, lambda: frames)

  return frames
//  [python]
def resize_smallest(frames: tf.Tensor, min_resize: int) -> tf.Tensor:
  """Resizes frames so that min(`height`, `width`) is equal to `min_resize`.

  This function will not do anything if the min(`height`, `width`) is already
  equal to `min_resize`. This allows to save compute time.

  Args:
    frames: A Tensor of dimension [timesteps, input_h, input_w, channels].
    min_resize: Minimum size of the final image dimensions.

  Returns:
    A Tensor of shape [timesteps, output_h, output_w, channels] of type
      frames.dtype where min(output_h, output_w) = min_resize.
  """
  shape = tf.shape(frames)
  input_h = shape[1]
  input_w = shape[2]

  output_h = tf.maximum(min_resize, (input_h * min_resize) // input_w)
  output_w = tf.maximum(min_resize, (input_w * min_resize) // input_h)

  def resize_fn():
    frames_resized = tf.image.resize(frames, (output_h, output_w))
    return tf.cast(frames_resized, frames.dtype)

  should_resize = tf.math.logical_or(
      tf.not_equal(input_w, output_w), tf.not_equal(input_h, output_h))
  frames = tf.cond(should_resize, resize_fn, lambda: frames)

  return frames
//  [python]
def sample_segment_sequence(sequence: tf.Tensor,
                            num_frames: int,
                            is_training: bool,
                            seed: Optional[int] = None) -> tf.Tensor:
  """Samples a single segment of size `num_frames` from a given sequence.

  This function follows the temporal segment network sampling style
  (https://arxiv.org/abs/1608.00859). The video sequence would be divided into
  `num_frames` non-overlapping segments with same length. If `is_training` is
  `True`, we would randomly sampling one frame for each segment, and when
  `is_training` is `False`, only the center frame of each segment is sampled.

  Args:
    sequence: Any tensor where the first dimension is timesteps.
    num_frames: Number of frames to take.
    is_training: A boolean indicating sampling in training or evaluation mode.
    seed: A deterministic seed to use when sampling.

  Returns:
    A single `tf.Tensor` with first dimension `num_steps` with the sampled
    segment.
  """
  sequence_length = tf.shape(sequence)[0]

  sequence_length = tf.to_float(sequence_length)
  segment_length = tf.to_float(sequence_length // num_frames)
  segment_indices = tf.linspace(0.0, sequence_length, num_frames + 1)
  segment_indices = tf.to_int32(segment_indices)

  if is_training:
    segment_length = tf.to_int32(segment_length)
    # pylint:disable=g-long-lambda
    segment_offsets = tf.cond(
        segment_length == 0,
        lambda: tf.zeros(shape=(num_frames,), dtype=tf.int32),
        lambda: tf.random.uniform(
            shape=(num_frames,),
            minval=0,
            maxval=segment_length,
            dtype=tf.int32,
            seed=seed))
    # pylint:disable=g-long-lambda

  else:
    # Only sampling central frame during inference for being deterministic.
    segment_offsets = tf.ones(
        shape=(num_frames,), dtype=tf.int32) * tf.to_int32(
            segment_length // 2)

  indices = segment_indices[:-1] + segment_offsets
  indices.set_shape((num_frames,))

  return tf.gather(sequence, indices)
//  [python]
def sample_segment_sequence(sequence: tf.Tensor,
                            num_frames: int,
                            is_training: bool,
                            seed: Optional[int] = None) -> tf.Tensor:
  """Samples a single segment of size `num_frames` from a given sequence.

  This function follows the temporal segment network sampling style
  (https://arxiv.org/abs/1608.00859). The video sequence would be divided into
  `num_frames` non-overlapping segments with same length. If `is_training` is
  `True`, we would randomly sampling one frame for each segment, and when
  `is_training` is `False`, only the center frame of each segment is sampled.

  Args:
    sequence: Any tensor where the first dimension is timesteps.
    num_frames: Number of frames to take.
    is_training: A boolean indicating sampling in training or evaluation mode.
    seed: A deterministic seed to use when sampling.

  Returns:
    A single `tf.Tensor` with first dimension `num_steps` with the sampled
    segment.
  """
  sequence_length = tf.shape(sequence)[0]

  sequence_length = tf.cast(sequence_length, tf.float32)
  segment_length = tf.cast(sequence_length // num_frames, tf.float32)
  segment_indices = tf.linspace(0.0, sequence_length, num_frames + 1)
  segment_indices = tf.cast(segment_indices, tf.int32)

  if is_training:
    segment_length = tf.cast(segment_length, tf.int32)
    # pylint:disable=g-long-lambda
    segment_offsets = tf.cond(
        segment_length == 0,
        lambda: tf.zeros(shape=(num_frames,), dtype=tf.int32),
        lambda: tf.random.uniform(
            shape=(num_frames,),
            minval=0,
            maxval=segment_length,
            dtype=tf.int32,
            seed=seed))
    # pylint:disable=g-long-lambda

  else:
    # Only sampling central frame during inference for being deterministic.
    segment_offsets = tf.ones(
        shape=(num_frames,), dtype=tf.int32) * tf.cast(
            segment_length // 2, dtype=tf.int32)

  indices = segment_indices[:-1] + segment_offsets
  indices.set_shape((num_frames,))

  return tf.gather(sequence, indices)
//  [python]
def create_variables(self):
        with tf.name_scope("model_inputs"):
            # raw state representation
            self.states = tf.placeholder(tf.float32, [None, self.max_layers*4], name="states")

        with tf.name_scope("predict_actions"):
            # initialize policy network
            with tf.variable_scope("policy_network"):
                self.policy_outputs = self.policy_network(self.states, self.max_layers)

            self.action_scores = tf.identity(self.policy_outputs, name="action_scores")

            self.predicted_action = tf.cast(tf.scalar_mul(self.division_rate, self.action_scores), tf.int32, name="predicted_action")


        # regularization loss
        policy_network_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="policy_network")

        # compute loss and gradients
        with tf.name_scope("compute_gradients"):
            # gradients for selecting action from policy network
            self.discounted_rewards = tf.placeholder(tf.float32, (None,), name="discounted_rewards")

            with tf.variable_scope("policy_network", reuse=True):
                self.logprobs = self.policy_network(self.states, self.max_layers)
                print("self.logprobs", self.logprobs)

            # compute policy loss and regularization loss
            self.cross_entropy_loss = tf.nn.softmax_cross_entropy_with_logits(logits=self.logprobs[:, -1, :], labels=self.states)
            self.pg_loss            = tf.reduce_mean(self.cross_entropy_loss)
            self.reg_loss           = tf.reduce_sum([tf.reduce_sum(tf.square(x)) for x in policy_network_variables]) # Regularization
            self.loss               = self.pg_loss + self.reg_param * self.reg_loss

            #compute gradients
            self.gradients = self.optimizer.compute_gradients(self.loss)
            
            # compute policy gradients
            for i, (grad, var) in enumerate(self.gradients):
                if grad is not None:
                    self.gradients[i] = (grad * self.discounted_rewards, var)

            # training update
            with tf.name_scope("train_policy_network"):
                # apply gradients to update policy network
                self.train_op = self.optimizer.apply_gradients(self.gradients, global_step=self.global_step)
//  [python]
def testTrainPerArmAgent(self, agent_class, use_num_samples_list):
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(
        2, 3, 4, add_num_actions_feature=True
    )
    time_step_spec = ts.time_step_spec(obs_spec)
    reward_net = (
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
    )
    optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=0.1)
    if use_num_samples_list:
      agent_class = functools.partial(
          agent_class, num_samples_list=self._per_arm_num_samples_list
      )
    agent = agent_class(
        time_step_spec,
        self._action_spec,
        reward_network=reward_net,
        accepts_per_arm_features=True,
        optimizer=optimizer,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: tf.cast(
            tf.reshape(tf.range(24), shape=[2, 4, 3]), dtype=tf.float32
        ),
        bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY: tf.ones([2], dtype=tf.int32),
    }
    actions = np.array([0, 3], dtype=np.int32)
    rewards = np.array([0.5, 3.0], dtype=np.float32)
    initial_step, final_step = _get_initial_and_final_steps(
        observations, rewards
    )
    action_step = policy_step.PolicyStep(
        action=tf.convert_to_tensor(actions),
        info=policy_utilities.PerArmPolicyInfo(
            chosen_arm_features=np.array(
                [[1, 2, 3], [3, 2, 1]], dtype=np.float32
            )
        ),
    )
    experience = _get_experience(initial_step, action_step, final_step)
    self.evaluate(tf.compat.v1.initialize_all_variables())
    self.evaluate(agent.train(experience, None).loss)
    if use_num_samples_list:
      self.evaluate(agent.train(experience, None).loss)
      self.assertEqual(
          self.evaluate(self._per_arm_num_samples_list[0].read_value()), 4
      )
//  [python]
def pad_to_fixed_size(input_tensor, size, constant_values=0):
  """Pads data to a fixed length at the first dimension.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.
    constant_values: `int` value assigned to the paddings.

  Returns:
    `Tensor` with the first dimension padded to `size`.
  """
  input_shape = input_tensor.get_shape().as_list()
  padding_shape = []

  # Computes the padding length on the first dimension.
  padding_length = tf.maximum(0, size - tf.shape(input_tensor)[0])
  assert_length = tf.Assert(
      tf.greater_equal(padding_length, 0), [padding_length])
  with tf.control_dependencies([assert_length]):
    padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input=input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = t.cast(constant_values * tf.ones(padding_shape),
                     input_tensor.dtype)
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  output_shape = input_shape
  output_shape[0] = size
  padded_tensor.set_shape(output_shape)
  return padded_tensor
//  [python]
def pad_to_fixed_size(input_tensor, size, constant_values=0):
  """Pads data to a fixed length at the first dimension.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.
    constant_values: `int` value assigned to the paddings.

  Returns:
    `Tensor` with the first dimension padded to `size`.
  """
  input_shape = input_tensor.get_shape().as_list()
  padding_shape = []

  # Computes the padding length on the first dimension.
  padding_length = tf.maximum(0, size - tf.shape(input_tensor)[0])
  assert_length = tf.Assert(
      tf.greater_equal(padding_length, 0), [padding_length])
  with tf.control_dependencies([assert_length]):
    padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input=input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = tf.cast(constant_values * tf.ones(padding_shape),
                     input_tensor.dtype)
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  output_shape = input_shape
  output_shape[0] = size
  padded_tensor.set_shape(output_shape)
  return padded_tensor
//  [python]
def resize_and_crop_masks(masks, image_scale, output_size, offset):
  """Resizes boxes to output size with scale and offset.

  Args:
    masks: `Tensor` of shape [N, H, W, 1] representing ground truth masks.
    image_scale: 2D float `Tensor` representing scale factors that apply to
      [height, width] of input image.
    output_size: 2D `Tensor` or `int` representing [height, width] of target
      output image size.
    offset: 2D `Tensor` representing top-left corner [y0, x0] to crop scaled
      boxes.

  Returns:
    masks: `Tensor` of shape [N, H, W, 1] representing the scaled masks.
  """
  mask_size = tf.shape(input=masks)[1:3]
  scaled_size = tf.to_int32(image_scale * tf.to_float(mask_size))
  scaled_masks = tf.image.resize(
      masks, scaled_size, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  offset = tf.to_int32(offset)
  scaled_masks = scaled_masks[:, offset[0]:offset[0] + output_size[0],
                              offset[1]:offset[1] + output_size[1], :]

  output_masks = tf.image.pad_to_bounding_box(scaled_masks, 0, 0,
                                              output_size[0], output_size[1])
  return output_masks
//  [python]
def resize_and_crop_masks(masks, image_scale, output_size, offset):
  """Resizes boxes to output size with scale and offset.

  Args:
    masks: `Tensor` of shape [N, H, W, 1] representing ground truth masks.
    image_scale: 2D float `Tensor` representing scale factors that apply to
      [height, width] of input image.
    output_size: 2D `Tensor` or `int` representing [height, width] of target
      output image size.
    offset: 2D `Tensor` representing top-left corner [y0, x0] to crop scaled
      boxes.

  Returns:
    masks: `Tensor` of shape [N, H, W, 1] representing the scaled masks.
  """
  mask_size = tf.shape(input=masks)[1:3]
  scaled_size = tf.cast(image_scale * tf.cast(mask_size, image_scale.dtype),
                        tf.int32)
  scaled_masks = tf.image.resize(
      masks, scaled_size, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  offset = tf.cast(offset, tf.int32)
  scaled_masks = scaled_masks[:, offset[0]:offset[0] + output_size[0],
                              offset[1]:offset[1] + output_size[1], :]

  output_masks = tf.image.pad_to_bounding_box(scaled_masks, 0, 0,
                                              output_size[0], output_size[1])
  return output_masks
//  [python]
def sparse_softmax_cross_entropy_with_logits(y, y_hat, label_smoothing=0.0):
    label_smoothing = tf.cast(tf.convert_to_tensor(label_smoothing), y_hat.dtype)
    n_classes = tf.shape(y_hat)[-1]
    y = tf.one_hot(tf.cast(y, tf.int32), n_classes, dtype=y_hat.dtype)
    y = y * (1. - label_smoothing) + (label_smoothing / tf.cast(n_classes, y_hat.dtype))
    return tf.nn.softmax_cross_entropy_with_logits(y, y_hat)
//  [python]
def create_dataset_from_data_producer(producer, params):
  """Return dataset online-generating data."""

  def preprocess_train_input(features, labels):
    """Pre-process the training data.

    This is needed because
    - The label needs to be extended to be used in the loss fn
    - We need the same inputs for training and eval so adding fake inputs
      for DUPLICATE_MASK in training data.

    Args:
      features: Dictionary of features for training.
      labels: Training labels.

    Returns:
      Processed training features.
    """
    fake_dup_mask = tf.zeros_like(features[movielens.USER_COLUMN])
    features[rconst.DUPLICATE_MASK] = fake_dup_mask
    features[rconst.TRAIN_LABEL_KEY] = labels
    return features

  train_input_fn = producer.make_input_fn(is_training=True)
  train_input_dataset = train_input_fn(params).map(preprocess_train_input)

  def preprocess_eval_input(features):
    """Pre-process the eval data.

    This is needed because:
    - The label needs to be extended to be used in the loss fn
    - We need the same inputs for training and eval so adding fake inputs
      for VALID_PT_MASK in eval data.

    Args:
      features: Dictionary of features for evaluation.

    Returns:
      Processed evaluation features.
    """
    labels = tf.to_int32(tf.zeros_like(features[movielens.USER_COLUMN]))
    fake_valid_pt_mask = tf.to_int32(
        tf.zeros_like(features[movielens.USER_COLUMN]))
    features[rconst.VALID_POINT_MASK] = fake_valid_pt_mask
    features[rconst.TRAIN_LABEL_KEY] = labels
    return features

  eval_input_fn = producer.make_input_fn(is_training=False)
  eval_input_dataset = eval_input_fn(params).map(preprocess_eval_input)

  return train_input_dataset, eval_input_dataset