jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def img2mpi(self, img, depth, planedepths):
    """Compute ground truth MPI of visible content using depth map."""

    height = tf.shape(img)[1]
    width = tf.shape(img)[2]
    num_depths = planedepths.shape[0]
    depth_inds = (tf.to_float(num_depths) - 1) * (
        (1.0 / depth) - (1.0 / planedepths[0])) / ((1.0 / planedepths[-1]) -
                                                   (1.0 / planedepths[0]))
    depth_inds = tf.round(depth_inds)
    depth_inds_tile = tf.to_int32(
        tf.tile(depth_inds[:, :, :, tf.newaxis], [1, 1, 1, num_depths]))
    _, _, d = tf.meshgrid(
        tf.range(height), tf.range(width), tf.range(num_depths), indexing='ij')
    mpi_colors = tf.to_float(
        tf.tile(img[:, :, :, tf.newaxis, :], [1, 1, 1, num_depths, 1]))
    mpi_alphas = tf.to_float(
        tf.where(
            tf.equal(depth_inds_tile, d), tf.ones_like(depth_inds_tile),
            tf.zeros_like(depth_inds_tile)))
    mpi = tf.concat([mpi_colors, mpi_alphas[Ellipsis, tf.newaxis]], axis=4)
    return mpi
//  [python]
def img2mpi(self, img, depth, planedepths):
    """Compute ground truth MPI of visible content using depth map."""

    height = tf.shape(img)[1]
    width = tf.shape(img)[2]
    num_depths = planedepths.shape[0]
    depth_inds = (tf.to_float(num_depths) - 1) * (
        (1.0 / depth) - (1.0 / planedepths[0])) / ((1.0 / planedepths[-1]) -
                                                   (1.0 / planedepths[0]))
    depth_inds = tf.round(depth_inds)
    depth_inds_tile = tf.cast(
        tf.tile(depth_inds[:, :, :, tf.newaxis], [1, 1, 1, num_depths]), tf.int32)
    _, _, d = tf.meshgrid(
        tf.range(height), tf.range(width), tf.range(num_depths), indexing='ij')
    mpi_colors = tf.to_float(
        tf.tile(img[:, :, :, tf.newaxis, :], [1, 1, 1, num_depths, 1]))
    mpi_alphas = tf.to_float(
        tf.where(
            tf.equal(depth_inds_tile, d), tf.ones_like(depth_inds_tile),
            tf.zeros_like(depth_inds_tile)))
    mpi = tf.concat([mpi_colors, mpi_alphas[Ellipsis, tf.newaxis]], axis=4)
    return mpi
//  [python]
def _resource_apply_dense(self, grad, handle):
    var = handle
    grad = tf.to_float(grad)
    grad_squared = tf.square(grad) + self._epsilon1
    grad_squared_mean = tf.reduce_mean(grad_squared)
    decay_rate = self._decay_rate
    update_scale = self._learning_rate
    old_val = var
    if var.dtype.base_dtype == tf.bfloat16:
      old_val = tf.to_float(self._parameter_encoding.decode(old_val))
    if self._multiply_by_parameter_scale:
      update_scale *= tf.to_float(self._parameter_scale(old_val))
    # HACK: Make things dependent on grad.
    # This confounds the XLA rewriter and keeps it from fusing computations
    # across different variables.  This fusion is a bad for HBM usage, since
    # it causes the gradients to persist in memory.
    decay_rate += grad_squared_mean * 1e-30
    update_scale += grad_squared_mean * 1e-30
    # END HACK
    mixing_rate = 1.0 - decay_rate
    shape = var.get_shape().as_list()
    updates = []
    if self._should_use_factored_second_moment_estimate(shape):
      grad_squared_row_mean = tf.reduce_mean(grad_squared, -1)
      grad_squared_col_mean = tf.reduce_mean(grad_squared, -2)
      vr = self.get_slot(var, "vr")
      new_vr = (decay_rate * vr + mixing_rate * grad_squared_row_mean)
      vc = self.get_slot(var, "vc")
      new_vc = (decay_rate * vc + mixing_rate * grad_squared_col_mean)
      vr_update = tf.assign(vr, new_vr, use_locking=self._use_locking)
      vc_update = tf.assign(vc, new_vc, use_locking=self._use_locking)
      updates = [vr_update, vc_update]
      long_term_mean = tf.reduce_mean(new_vr, -1, keepdims=True)
      r_factor = tf.rsqrt(new_vr / long_term_mean)
      c_factor = tf.rsqrt(new_vc)
      x = grad * tf.expand_dims(r_factor, -1) * tf.expand_dims(c_factor, -2)
    else:
      v = self.get_slot(var, "v")
      new_v = decay_rate * v + mixing_rate * grad_squared
      v_update = tf.assign(v, new_v, use_locking=self._use_locking)
      updates = [v_update]
      x = grad * tf.rsqrt(new_v)
    if self._clipping_threshold is not None:
      clipping_denom = tf.maximum(1.0, reduce_rms(x) / self._clipping_threshold)
      x /= clipping_denom
    subtrahend = update_scale * x
    if self._beta1:
      m = self.get_slot(var, "m")
      new_m = self._beta1 * tf.to_float(m) + (1.0 - self._beta1) * subtrahend
      subtrahend = new_m
      new_m = common_layers.cast_like(new_m, var)
      updates.append(tf.assign(m, new_m, use_locking=self._use_locking))
    new_val = tf.to_float(old_val) - subtrahend
    if var.dtype.base_dtype == tf.bfloat16:
      new_val = self._parameter_encoding.encode(
          new_val, self._quantization_noise)
    if self._simulated_quantize_bits:
      new_val = quantization.simulated_quantize(
          var - subtrahend, self._simulated_quantize_bits,
          self._quantization_noise)
    var_update = tf.assign(var, new_val, use_locking=self._use_locking)
    updates = [var_update] + updates
    return tf.group(*updates)
//  [python]
def _resource_apply_dense(self, grad, handle):
    var = handle
    grad = tf.cast(grad, tf.float32)
    grad_squared = tf.square(grad) + self._epsilon1
    grad_squared_mean = tf.reduce_mean(grad_squared)
    decay_rate = self._decay_rate
    update_scale = self._learning_rate
    old_val = var
    if var.dtype.base_dtype == tf.bfloat16:
      old_val = tf.cast(self._parameter_encoding.decode(old_val), tf.float32)
    if self._multiply_by_parameter_scale:
      update_scale *= tf.cast(self._parameter_scale(old_val), tf.float32)
    # HACK: Make things dependent on grad.
    # This confounds the XLA rewriter and keeps it from fusing computations
    # across different variables.  This fusion is a bad for HBM usage, since
    # it causes the gradients to persist in memory.
    decay_rate += grad_squared_mean * 1e-30
    update_scale += grad_squared_mean * 1e-30
    # END HACK
    mixing_rate = 1.0 - decay_rate
    shape = var.get_shape().as_list()
    updates = []
    if self._should_use_factored_second_moment_estimate(shape):
      grad_squared_row_mean = tf.reduce_mean(grad_squared, -1)
      grad_squared_col_mean = tf.reduce_mean(grad_squared, -2)
      vr = self.get_slot(var, "vr")
      new_vr = (decay_rate * vr + mixing_rate * grad_squared_row_mean)
      vc = self.get_slot(var, "vc")
      new_vc = (decay_rate * vc + mixing_rate * grad_squared_col_mean)
      vr_update = tf.assign(vr, new_vr, use_locking=self._use_locking)
      vc_update = tf.assign(vc, new_vc, use_locking=self._use_locking)
      updates = [vr_update, vc_update]
      long_term_mean = tf.reduce_mean(new_vr, -1, keepdims=True)
      r_factor = tf.rsqrt(new_vr / long_term_mean)
      c_factor = tf.rsqrt(new_vc)
      x = grad * tf.expand_dims(r_factor, -1) * tf.expand_dims(c_factor, -2)
    else:
      v = self.get_slot(var, "v")
      new_v = decay_rate * v + mixing_rate * grad_squared
      v_update = tf.assign(v, new_v, use_locking=self._use_locking)
      updates = [v_update]
      x = grad * tf.rsqrt(new_v)
    if self._clipping_threshold is not None:
      clipping_denom = tf.maximum(1.0, reduce_rms(x) / self._clipping_threshold)
      x /= clipping_denom
    subtrahend = update_scale * x
    if self._beta1:
      m = self.get_slot(var, "m")
      new_m = self._beta1 * tf.cast(m, tf.float32) + (1.0 - self._beta1) * subtrahend
      subtrahend = new_m
      new_m = common_layers.cast_like(new_m, var)
      updates.append(tf.assign(m, new_m, use_locking=self._use_locking))
    new_val = tf.cast(old_val, tf.float32) - subtrahend
    if var.dtype.base_dtype == tf.bfloat16:
      new_val = self._parameter_encoding.encode(
          new_val, self._quantization_noise)
    if self._simulated_quantize_bits:
      new_val = quantization.simulated_quantize(
          var - subtrahend, self._simulated_quantize_bits,
          self._quantization_noise)
    var_update = tf.assign(var, new_val, use_locking=self._use_locking)
    updates = [var_update] + updates
    return tf.group(*updates)
//  [python]
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
//  [python]
def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
        t = example[name]
        if t.dtype == tf.int64:
            t = tf.cast(t, tf.int32)
        example[name] = t
    return example
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']
        seqs_units = self.item_embs(seqs_id)

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

        # Recurrency
        with tf.variable_scope("S2PNM/Reccurency"):
            h, _ = self._cudnn_rnn(tf.transpose(seqs_units, [1, 0, 2]))
            h = tf.transpose(h, [1, 0, 2])  # mask the hidden states as dynamic_rnn did

        seqs_outs = h * seqs_masks

        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_units = self.item_embs(seqs_id)

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

        # Recurrency
        with tf.variable_scope("S2PNM/Reccurency"):
            h, _ = self._cudnn_rnn(tf.transpose(seqs_units, [1, 0, 2]))
            h = tf.transpose(h, [1, 0, 2])  # mask the hidden states as dynamic_rnn did

        seqs_outs = h * seqs_masks

        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 compute_loss(self, decoder_output, _features, labels):
    """Computes the loss for this model.

    Returns a tuple `(losses, loss)`, where `losses` are the per-batch
    losses and loss is a single scalar tensor to minimize.
    """
    #pylint: disable=R0201
    # Calculate loss per example-timestep of shape [B, T]
    losses = seq2seq_losses.cross_entropy_sequence_loss(
        logits=decoder_output.logits[:, :, :],
        targets=tf.transpose(labels["target_ids"][:, 1:], [1, 0]),
        sequence_length=labels["target_len"] - 1)

    # Calculate the average log perplexity
    loss = tf.reduce_sum(losses) / tf.to_float(
        tf.reduce_sum(labels["target_len"] - 1))

    return losses, loss
//  [python]
def compute_loss(self, decoder_output, _features, labels):
    """Computes the loss for this model.

    Returns a tuple `(losses, loss)`, where `losses` are the per-batch
    losses and loss is a single scalar tensor to minimize.
    """
    #pylint: disable=R0201
    # Calculate loss per example-timestep of shape [B, T]
    losses = seq2seq_losses.cross_entropy_sequence_loss(
        logits=decoder_output.logits[:, :, :],
        targets=tf.transpose(labels["target_ids"][:, 1:], [1, 0]),
        sequence_length=labels["target_len"] - 1)

    # Calculate the average log perplexity
    loss = tf.reduce_sum(losses) / tf.cast(
        tf.reduce_sum(labels["target_len"] - 1), tf.float32)

    return losses, loss
//  [python]
def _create_predictions(self, decoder_output, features, labels, losses=None):
    """Creates the dictionary of predictions that is returned by the model.
    """
    predictions = {}

    # Add features and, if available, labels to predictions
    predictions.update(_flatten_dict({"features": features}))
    if labels is not None:
      predictions.update(_flatten_dict({"labels": labels}))

    if losses is not None:
      predictions["losses"] = _transpose_batch_time(losses)

    # Decoders returns output in time-major form [T, B, ...]
    # Here we transpose everything back to batch-major for the user
    output_dict = collections.OrderedDict(
        zip(decoder_output._fields, decoder_output))
    decoder_output_flat = _flatten_dict(output_dict)
    decoder_output_flat = {
        k: _transpose_batch_time(v)
        for k, v in decoder_output_flat.items()
    }
    predictions.update(decoder_output_flat)

    # If we predict the ids also map them back into the vocab and process them
    if "predicted_ids" in predictions.keys():
      vocab_tables = graph_utils.get_dict_from_collection("vocab_tables")
      target_id_to_vocab = vocab_tables["target_id_to_vocab"]
      predicted_tokens = target_id_to_vocab.lookup(
          tf.to_int64(predictions["predicted_ids"]))
      # Raw predicted tokens
      predictions["predicted_tokens"] = predicted_tokens

    return predictions
//  [python]
def _create_predictions(self, decoder_output, features, labels, losses=None):
    """Creates the dictionary of predictions that is returned by the model.
    """
    predictions = {}

    # Add features and, if available, labels to predictions
    predictions.update(_flatten_dict({"features": features}))
    if labels is not None:
      predictions.update(_flatten_dict({"labels": labels}))

    if losses is not None:
      predictions["losses"] = _transpose_batch_time(losses)

    # Decoders returns output in time-major form [T, B, ...]
    # Here we transpose everything back to batch-major for the user
    output_dict = collections.OrderedDict(
        zip(decoder_output._fields, decoder_output))
    decoder_output_flat = _flatten_dict(output_dict)
    decoder_output_flat = {
        k: _transpose_batch_time(v)
        for k, v in decoder_output_flat.items()
    }
    predictions.update(decoder_output_flat)

    # If we predict the ids also map them back into the vocab and process them
    if "predicted_ids" in predictions.keys():
      vocab_tables = graph_utils.get_dict_from_collection("vocab_tables")
      target_id_to_vocab = vocab_tables["target_id_to_vocab"]
      predicted_tokens = target_id_to_vocab.lookup(
          tf.cast(predictions["predicted_ids"], tf.int64))
      # Raw predicted tokens
      predictions["predicted_tokens"] = predicted_tokens

    return predictions
//  [python]
def get_position_signal(sequence_length, position_dim=8):
  """Return fixed position signal as sine waves.

  Sine waves frequencies are linearly spaced so that shortest is 2 and
  longest is half the maximum length. That way the longest frequency
  is long enough to be monotonous over the whole sequence length.
  Sine waves are also shifted so that they don't all start with the same
  value.
  We don't use learned positional embeddings because these embeddings are
  projected linearly along with the original embeddings, and the projection is
  learned.

  Args:
    sequence_length: int, T, length of the sequence..
    position_dim: int, P, number of sine waves.

  Returns:
    A [T, P] tensor, position embeddings.
  """
  # Compute the frequencies.
  periods = tf.exp(
      tf.lin_space(
          tf.log(2.0), tf.log(tf.to_float(sequence_length)), position_dim))
  frequencies = 1.0 / periods  # Shape [T, P].

  # Compute the sine waves.
  xs = frequencies[None, :] * tf.to_float(tf.range(sequence_length)[:, None])
  shifts = tf.lin_space(0.0, 2.0, position_dim)[None, :]  # [1, P]
  positions = tf.math.cos(math.pi * (xs + shifts))  # [T, P]
  positions.shape.assert_is_compatible_with([sequence_length, position_dim])
  return positions
//  [python]
def get_position_signal(sequence_length, position_dim=8):
  """Return fixed position signal as sine waves.

  Sine waves frequencies are linearly spaced so that shortest is 2 and
  longest is half the maximum length. That way the longest frequency
  is long enough to be monotonous over the whole sequence length.
  Sine waves are also shifted so that they don't all start with the same
  value.
  We don't use learned positional embeddings because these embeddings are
  projected linearly along with the original embeddings, and the projection is
  learned.

  Args:
    sequence_length: int, T, length of the sequence..
    position_dim: int, P, number of sine waves.

  Returns:
    A [T, P] tensor, position embeddings.
  """
  # Compute the frequencies.
  periods = tf.exp(
      tf.lin_space(
          tf.log(2.0), tf.log(tf.cast(sequence_length, tf.float32)), position_dim))
  frequencies = 1.0 / periods  # Shape [T, P].

  # Compute the sine waves.
  xs = frequencies[None, :] * tf.cast(tf.range(sequence_length)[:, None], tf.float32)
  shifts = tf.lin_space(0.0, 2.0, position_dim)[None, :]  # [1, P]
  positions = tf.math.cos(math.pi * (xs + shifts))  # [T, P]
  positions.shape.assert_is_compatible_with([sequence_length, position_dim])
  return positions
//  [python]
def get_learning_rate_decay(learning_rate, global_step, params):
    if params.learning_rate_decay in ["linear_warmup_rsqrt_decay", "noam"]:
        step = tf.to_float(global_step)
        warmup_steps = tf.to_float(params.warmup_steps)
        multiplier = params.hidden_size ** -0.5
        decay = multiplier * tf.minimum((step + 1) * (warmup_steps ** -1.5),
                                        (step + 1) ** -0.5)

        return learning_rate * decay
    elif params.learning_rate_decay == "piecewise_constant":
        return tf.train.piecewise_constant(tf.to_int32(global_step),
                                           params.learning_rate_boundaries,
                                           params.learning_rate_values)
    elif params.learning_rate_decay == "none":
        return learning_rate
    else:
        raise ValueError("Unknown learning_rate_decay")
//  [python]
def get_learning_rate_decay(learning_rate, global_step, params):
    if params.learning_rate_decay in ["linear_warmup_rsqrt_decay", "noam"]:
        step = tf.cast(global_step, tf.float32)
        warmup_steps = tf.cast(params.warmup_steps, tf.float32)
        multiplier = params.hidden_size ** -0.5
        decay = multiplier * tf.minimum((step + 1) * (warmup_steps ** -1.5),
                                        (step + 1) ** -0.5)

        return learning_rate * decay
    elif params.learning_rate_decay == "piecewise_constant":
        return tf.train.piecewise_constant(tf.cast(global_step, tf.int32),
                                           params.learning_rate_boundaries,
                                           params.learning_rate_values)
    elif params.learning_rate_decay == "none":
        return learning_rate
    else:
        raise ValueError("Unknown learning_rate_decay")
//  [python]
def format_network_input(self, ref_image, psv_src_images, ref_pose,
                           psv_src_poses, planes, intrinsics):
    """Format the network input.

    Args:
      ref_image: reference source image [batch, height, width, 3]
      psv_src_images: stack of source images (excluding the ref image) [batch,
        height, width, 3*(#source)]
      ref_pose: reference camera-to-world pose (where PSV is constructed)
        [batch, 4, 4]
      psv_src_poses: input poses (camera to world) [batch, 4, 4, #source]
      planes: list of scalar depth values for each plane
      intrinsics: camera intrinsics [batch, 3, 3]

    Returns:
      net_input: [batch, height, width, #planes, (#source+1)*3]
    """

    batch_size = tf.shape(psv_src_images)[0]
    height = tf.shape(psv_src_images)[1]
    width = tf.shape(psv_src_images)[2]
    _, _, _, num_psv_source = psv_src_poses.get_shape().as_list()
    num_planes = tf.shape(planes)[0]

    filler = tf.concat(
        [tf.zeros([batch_size, 1, 3]),
         tf.ones([batch_size, 1, 1])], axis=2)
    intrinsics_filler = tf.stack([
        tf.to_float(height),
        tf.to_float(width),
        tf.to_float(intrinsics[0, 0, 0])
    ],
                                 axis=0)[:, tf.newaxis]

    ref_pose_c2w = ref_pose
    ref_pose_c2w = tf.concat([
        tf.concat([
            ref_pose_c2w[:, :3, 0:1], ref_pose_c2w[:, :3, 1:2],
            -1.0 * ref_pose_c2w[:, :3, 2:3], ref_pose_c2w[:, :3, 3:]
        ],
                  axis=2), filler
    ],
                             axis=1)
    ref_pose_c2w = tf.concat([ref_pose_c2w[0, :3, :], intrinsics_filler],
                             axis=1)

    net_input = []
    for i in range(num_psv_source):
      curr_pose_c2w = psv_src_poses[:, :, :, i]
      curr_pose_c2w = tf.concat([
          tf.concat([
              curr_pose_c2w[:, :3, 0:1], curr_pose_c2w[:, :3, 1:2],
              -1.0 * curr_pose_c2w[:, :3, 2:3], curr_pose_c2w[:, :3, 3:]
          ], 2), filler
      ], 1)
      curr_pose_c2w = tf.concat([curr_pose_c2w[0, :3, :], intrinsics_filler],
                                axis=1)
      curr_image = psv_src_images[:, :, :, i * 3:(i + 1) * 3]
      curr_psv = pj.make_psv_homogs(curr_image, curr_pose_c2w, ref_pose_c2w,
                                    1.0 / planes, num_planes)
      net_input.append(curr_psv[tf.newaxis, Ellipsis])

    net_input = tf.concat(net_input, axis=4)
    ref_img_stack = tf.tile(
        tf.expand_dims(ref_image, 3), [1, 1, 1, num_planes, 1])
    net_input = tf.concat([ref_img_stack, net_input], axis=4)
    net_input.set_shape([1, None, None, None, 3 * (num_psv_source + 1)])

    return net_input
//  [python]
def format_network_input(self, ref_image, psv_src_images, ref_pose,
                           psv_src_poses, planes, intrinsics):
    """Format the network input.

    Args:
      ref_image: reference source image [batch, height, width, 3]
      psv_src_images: stack of source images (excluding the ref image) [batch,
        height, width, 3*(#source)]
      ref_pose: reference camera-to-world pose (where PSV is constructed)
        [batch, 4, 4]
      psv_src_poses: input poses (camera to world) [batch, 4, 4, #source]
      planes: list of scalar depth values for each plane
      intrinsics: camera intrinsics [batch, 3, 3]

    Returns:
      net_input: [batch, height, width, #planes, (#source+1)*3]
    """

    batch_size = tf.shape(psv_src_images)[0]
    height = tf.shape(psv_src_images)[1]
    width = tf.shape(psv_src_images)[2]
    _, _, _, num_psv_source = psv_src_poses.get_shape().as_list()
    num_planes = tf.shape(planes)[0]

    filler = tf.concat(
        [tf.zeros([batch_size, 1, 3]),
         tf.ones([batch_size, 1, 1])], axis=2)
    intrinsics_filler = tf.stack([
        tf.cast(height, tf.float32),
        tf.cast(width, tf.float32),
        tf.cast(intrinsics[0, 0, 0], tf.float32)
    ],
                                 axis=0)[:, tf.newaxis]

    ref_pose_c2w = ref_pose
    ref_pose_c2w = tf.concat([
        tf.concat([
            ref_pose_c2w[:, :3, 0:1], ref_pose_c2w[:, :3, 1:2],
            -1.0 * ref_pose_c2w[:, :3, 2:3], ref_pose_c2w[:, :3, 3:]
        ],
                  axis=2), filler
    ],
                             axis=1)
    ref_pose_c2w = tf.concat([ref_pose_c2w[0, :3, :], intrinsics_filler],
                             axis=1)

    net_input = []
    for i in range(num_psv_source):
      curr_pose_c2w = psv_src_poses[:, :, :, i]
      curr_pose_c2w = tf.concat([
          tf.concat([
              curr_pose_c2w[:, :3, 0:1], curr_pose_c2w[:, :3, 1:2],
              -1.0 * curr_pose_c2w[:, :3, 2:3], curr_pose_c2w[:, :3, 3:]
          ], 2), filler
      ], 1)
      curr_pose_c2w = tf.concat([curr_pose_c2w[0, :3, :], intrinsics_filler],
                                axis=1)
      curr_image = psv_src_images[:, :, :, i * 3:(i + 1) * 3]
      curr_psv = pj.make_psv_homogs(curr_image, curr_pose_c2w, ref_pose_c2w,
                                    1.0 / planes, num_planes)
      net_input.append(curr_psv[tf.newaxis, Ellipsis])

    net_input = tf.concat(net_input, axis=4)
    ref_img_stack = tf.tile(
        tf.expand_dims(ref_image, 3), [1, 1, 1, num_planes, 1])
    net_input = tf.concat([ref_img_stack, net_input], axis=4)
    net_input.set_shape([1, None, None, None, 3 * (num_psv_source + 1)])

    return net_input
//  [python]
def _decode_png_instance_masks(self, keys_to_tensors):
    """Decode PNG instance segmentation masks and stack into dense tensor.

    The instance segmentation masks are reshaped to [num_instances, height,
    width].

    Args:
      keys_to_tensors: a dictionary from keys to tensors.

    Returns:
      A 3-D float tensor of shape [num_instances, height, width] with values
        in {0, 1}.
    """

    def decode_png_mask(image_buffer):
      image = tf.squeeze(
          tf.image.decode_image(image_buffer, channels=1), axis=2)
      image.set_shape([None, None])
      image = tf.to_float(tf.greater(image, 0))
      return image

    png_masks = keys_to_tensors['image/object/mask']
    height = keys_to_tensors['image/height']
    width = keys_to_tensors['image/width']
    if isinstance(png_masks, tf.SparseTensor):
      png_masks = tf.sparse_tensor_to_dense(png_masks, default_value='')
    return tf.cond(
        tf.greater(tf.size(png_masks), 0),
        lambda: tf.map_fn(decode_png_mask, png_masks, dtype=tf.float32),
        lambda: tf.zeros(tf.to_int32(tf.stack([0, height, width]))))
//  [python]
def _decode_png_instance_masks(self, keys_to_tensors):
    """Decode PNG instance segmentation masks and stack into dense tensor.

    The instance segmentation masks are reshaped to [num_instances, height,
    width].

    Args:
      keys_to_tensors: a dictionary from keys to tensors.

    Returns:
      A 3-D float tensor of shape [num_instances, height, width] with values
        in {0, 1}.
    """

    def decode_png_mask(image_buffer):
      image = tf.squeeze(
          tf.image.decode_image(image_buffer, channels=1), axis=2)
      image.set_shape([None, None])
      image = tf.cast(tf.greater(image, 0), tf.float32)
      return image

    png_masks = keys_to_tensors['image/object/mask']
    height = keys_to_tensors['image/height']
    width = keys_to_tensors['image/width']
    if isinstance(png_masks, tf.SparseTensor):
      png_masks = tf.sparse_tensor_to_dense(png_masks, default_value='')
    return tf.cond(
        tf.greater(tf.size(png_masks), 0),
        lambda: tf.map_fn(decode_png_mask, png_masks, dtype=tf.float32),
        lambda: tf.zeros(tf.to_int32(tf.stack([0, height, width]))))
//  [python]
def mpi_render_view(self, input_mpi, ref_pose, tgt_pose, planes, intrinsics):
    """Render a target view from MPI representation.

    Args:
      input_mpi: input MPI [batch, height, width, #planes, 4]
      ref_pose: reference camera pose [batch, 4, 4]
      tgt_pose: target pose to render from [batch, 4, 4]
      planes: list of depths for each plane
      intrinsics: camera intrinsics [batch, 3, 3]

    Returns:
      rendered view [batch, height, width, 3]
    """

    batch_size, _, _ = tgt_pose.get_shape().as_list()
    num_planes = tf.shape(planes)[0]
    height = tf.shape(input_mpi)[1]
    width = tf.shape(input_mpi)[2]

    rgba_layers = input_mpi

    # render target viewpoint
    filler = tf.concat(
        [tf.zeros([batch_size, 1, 3]),
         tf.ones([batch_size, 1, 1])], axis=2)
    intrinsics_filler = tf.stack(
        [tf.to_float(height),
         tf.to_float(width), intrinsics[0, 0, 0]], axis=0)[:, tf.newaxis]

    ref_pose_c2w = ref_pose
    ref_pose_c2w = tf.concat([
        tf.concat([
            ref_pose_c2w[:, :3, 0:1], ref_pose_c2w[:, :3, 1:2],
            -1.0 * ref_pose_c2w[:, :3, 2:3], ref_pose_c2w[:, :3, 3:]
        ],
                  axis=2), filler
    ],
                             axis=1)
    ref_pose_c2w = tf.concat([ref_pose_c2w[0, :3, :], intrinsics_filler],
                             axis=1)

    tgt_pose_c2w = tgt_pose
    tgt_pose_c2w = tf.concat([
        tf.concat([
            tgt_pose_c2w[:, :3, 0:1], tgt_pose_c2w[:, :3, 1:2],
            -1.0 * tgt_pose_c2w[:, :3, 2:3], tgt_pose_c2w[:, :3, 3:]
        ],
                  axis=2), filler
    ],
                             axis=1)
    tgt_pose_c2w = tf.concat([tgt_pose_c2w[0, :3, :], intrinsics_filler],
                             axis=1)

    rendering, alpha_acc, accum = pj.render_mpi_homogs(
        rgba_layers,
        ref_pose_c2w,
        tgt_pose_c2w,
        1.0 / planes[0],
        1.0 / planes[-1],
        num_planes,
        debug=False)

    return rendering, alpha_acc, accum
//  [python]
def mpi_render_view(self, input_mpi, ref_pose, tgt_pose, planes, intrinsics):
    """Render a target view from MPI representation.

    Args:
      input_mpi: input MPI [batch, height, width, #planes, 4]
      ref_pose: reference camera pose [batch, 4, 4]
      tgt_pose: target pose to render from [batch, 4, 4]
      planes: list of depths for each plane
      intrinsics: camera intrinsics [batch, 3, 3]

    Returns:
      rendered view [batch, height, width, 3]
    """

    batch_size, _, _ = tgt_pose.get_shape().as_list()
    num_planes = tf.shape(planes)[0]
    height = tf.shape(input_mpi)[1]
    width = tf.shape(input_mpi)[2]

    rgba_layers = input_mpi

    # render target viewpoint
    filler = tf.concat(
        [tf.zeros([batch_size, 1, 3]),
         tf.ones([batch_size, 1, 1])], axis=2)
    intrinsics_filler = tf.stack(
        [tf.cast(height, tf.float32),
         tf.cast(width, tf.float32), intrinsics[0, 0, 0]], axis=0)[:, tf.newaxis]

    ref_pose_c2w = ref_pose
    ref_pose_c2w = tf.concat([
        tf.concat([
            ref_pose_c2w[:, :3, 0:1], ref_pose_c2w[:, :3, 1:2],
            -1.0 * ref_pose_c2w[:, :3, 2:3], ref_pose_c2w[:, :3, 3:]
        ],
                  axis=2), filler
    ],
                             axis=1)
    ref_pose_c2w = tf.concat([ref_pose_c2w[0, :3, :], intrinsics_filler],
                             axis=1)

    tgt_pose_c2w = tgt_pose
    tgt_pose_c2w = tf.concat([
        tf.concat([
            tgt_pose_c2w[:, :3, 0:1], tgt_pose_c2w[:, :3, 1:2],
            -1.0 * tgt_pose_c2w[:, :3, 2:3], tgt_pose_c2w[:, :3, 3:]
        ],
                  axis=2), filler
    ],
                             axis=1)
    tgt_pose_c2w = tf.concat([tgt_pose_c2w[0, :3, :], intrinsics_filler],
                             axis=1)

    rendering, alpha_acc, accum = pj.render_mpi_homogs(
        rgba_layers,
        ref_pose_c2w,
        tgt_pose_c2w,
        1.0 / planes[0],
        1.0 / planes[-1],
        num_planes,
        debug=False)

    return rendering, alpha_acc, accum
//  [python]
def _preprocess_zero_mean_unit_range(inputs):
  """Map image values from [0, 255] to [-1, 1]."""
  return (2.0 / 255.0) * tf.to_float(inputs) - 1.0
//  [python]
def _preprocess_zero_mean_unit_range(inputs):
  """Map image values from [0, 255] to [-1, 1]."""
  return (2.0 / 255.0) * tf.cast(inputs, tf.float32) - 1.0
//  [python]
def random_ccm():
  """Generates random RGB -> Camera color correction matrices."""
  # Takes a random convex combination of XYZ -> Camera CCMs.
  xyz2cams = [[[1.0234, -0.2969, -0.2266],
               [-0.5625, 1.6328, -0.0469],
               [-0.0703, 0.2188, 0.6406]],
              [[0.4913, -0.0541, -0.0202],
               [-0.613, 1.3513, 0.2906],
               [-0.1564, 0.2151, 0.7183]],
              [[0.838, -0.263, -0.0639],
               [-0.2887, 1.0725, 0.2496],
               [-0.0627, 0.1427, 0.5438]],
              [[0.6596, -0.2079, -0.0562],
               [-0.4782, 1.3016, 0.1933],
               [-0.097, 0.1581, 0.5181]]]
  num_ccms = len(xyz2cams)
  xyz2cams = tf.constant(xyz2cams)
  weights = tf.random_uniform((num_ccms, 1, 1), 1e-8, 1e8)
  weights_sum = tf.reduce_sum(weights, axis=0)
  xyz2cam = tf.reduce_sum(xyz2cams * weights, axis=0) / weights_sum

  # Multiplies with RGB -> XYZ to get RGB -> Camera CCM.
  rgb2xyz = tf.to_float([[0.4124564, 0.3575761, 0.1804375],
                         [0.2126729, 0.7151522, 0.0721750],
                         [0.0193339, 0.1191920, 0.9503041]])
  rgb2cam = tf.matmul(xyz2cam, rgb2xyz)

  # Normalizes each row.
  rgb2cam = rgb2cam / tf.reduce_sum(rgb2cam, axis=-1, keepdims=True)
  return rgb2cam
//  [python]
def random_ccm():
  """Generates random RGB -> Camera color correction matrices."""
  # Takes a random convex combination of XYZ -> Camera CCMs.
  xyz2cams = [[[1.0234, -0.2969, -0.2266],
               [-0.5625, 1.6328, -0.0469],
               [-0.0703, 0.2188, 0.6406]],
              [[0.4913, -0.0541, -0.0202],
               [-0.613, 1.3513, 0.2906],
               [-0.1564, 0.2151, 0.7183]],
              [[0.838, -0.263, -0.0639],
               [-0.2887, 1.0725, 0.2496],
               [-0.0627, 0.1427, 0.5438]],
              [[0.6596, -0.2079, -0.0562],
               [-0.4782, 1.3016, 0.1933],
               [-0.097, 0.1581, 0.5181]]]
  num_ccms = len(xyz2cams)
  xyz2cams = tf.constant(xyz2cams)
  weights = tf.random_uniform((num_ccms, 1, 1), 1e-8, 1e8)
  weights_sum = tf.reduce_sum(weights, axis=0)
  xyz2cam = tf.reduce_sum(xyz2cams * weights, axis=0) / weights_sum

  # Multiplies with RGB -> XYZ to get RGB -> Camera CCM.
  rgb2xyz = tf.cast([[0.4124564, 0.3575761, 0.1804375],
                         [0.2126729, 0.7151522, 0.0721750],
                         [0.0193339, 0.1191920, 0.9503041]], tf.float32)
  rgb2cam = tf.matmul(xyz2cam, rgb2xyz)

  # Normalizes each row.
  rgb2cam = rgb2cam / tf.reduce_sum(rgb2cam, axis=-1, keepdims=True)
  return rgb2cam
//  [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)
    name_to_features["answer_types"] = tf.FixedLenFeature([], tf.int64)

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

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

    return example

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

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

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

    return d

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

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

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

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

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

    return example

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

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

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

    return d

  return input_fn
//  [python]
def crop_image(x, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes."""
  original_shape = tf.shape(x)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(x), 3), ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.to_int32(tf.stack([offset_height, offset_width, 0]))

  with tf.control_dependencies([size_assertion]):
    x = tf.slice(x, offsets, cropped_shape)
    x = tf.reshape(x, cropped_shape)
  return x
//  [python]
def crop_image(x, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes."""
  original_shape = tf.shape(x)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(x), 3), ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.cast(tf.stack([offset_height, offset_width, 0]), tf.int32)

  with tf.control_dependencies([size_assertion]):
    x = tf.slice(x, offsets, cropped_shape)
    x = tf.reshape(x, cropped_shape)
  return x
//  [python]
def Binarize(image):
  d = tfd.Bernoulli(probs=image)
  return tf.to_float(d.sample())
//  [python]
def Binarize(image):
  d = tfd.Bernoulli(probs=image)
  return tf.cast(d.sample(), tf.float32)
//  [python]
def compute_pred_dict(candidates_dict, dev_features, raw_results):
  """Computes official answer key from raw logits."""
  raw_results_by_id = [(int(res["unique_id"] + 1), res) for res in raw_results]

  # Cast example id to int32 for each example, similarly to the raw results.
  sess = tf.Session()
  all_candidates = candidates_dict.items()
  example_ids = tf.to_int32(np.array([int(k) for k, _ in all_candidates
                                     ])).eval(session=sess)
  examples_by_id = zip(example_ids, all_candidates)

  # Cast unique_id also to int32 for features.
  feature_ids = []
  features = []
  for f in dev_features:
    feature_ids.append(f.features.feature["unique_ids"].int64_list.value[0] + 1)
    features.append(f.features.feature)
  feature_ids = tf.to_int32(np.array(feature_ids)).eval(session=sess)
  features_by_id = zip(feature_ids, features)

  # Join examples with features and raw results.
  examples = []
  merged = sorted(
      list(examples_by_id) + raw_results_by_id + list(features_by_id),
      key=lambda x: x[0])
  for idx, datum in merged:
    if isinstance(datum, tuple):
      examples.append(EvalExample(datum[0], datum[1]))
    elif "token_map" in datum:
      examples[-1].features[idx] = datum
    else:
      examples[-1].results[idx] = datum

  # Construct prediction objects.
  tf.logging.info("Computing predictions...")
  summary_dict = {}
  nq_pred_dict = {}
  for e in examples:
    summary = compute_predictions(e)
    summary_dict[e.example_id] = summary
    nq_pred_dict[e.example_id] = summary.predicted_label
    if len(nq_pred_dict) % 100 == 0:
      tf.logging.info("Examples processed: %d", len(nq_pred_dict))
  tf.logging.info("Done computing predictions.")

  return nq_pred_dict
//  [python]
def compute_pred_dict(candidates_dict, dev_features, raw_results):
  """Computes official answer key from raw logits."""
  raw_results_by_id = [(int(res["unique_id"] + 1), res) for res in raw_results]

  # Cast example id to int32 for each example, similarly to the raw results.
  sess = tf.compat.v1.Session()
  all_candidates = candidates_dict.items()
  example_ids = tf.cast(np.array([int(k) for k, _ in all_candidates]), tf.int32).eval(session=sess)
  examples_by_id = zip(example_ids, all_candidates)

  # Cast unique_id also to int32 for features.
  feature_ids = []
  features = []
  for f in dev_features:
    feature_ids.append(f.features.feature["unique_ids"].int64_list.value[0] + 1)
    features.append(f.features.feature)
  feature_ids = tf.cast(np.array(feature_ids), tf.int32).eval(session=sess)
  features_by_id = zip(feature_ids, features)

  # Join examples with features and raw results.
  examples = []
  merged = sorted(
      list(examples_by_id) + raw_results_by_id + list(features_by_id),
      key=lambda x: x[0])
  for idx, datum in merged:
    if isinstance(datum, tuple):
      examples.append(EvalExample(datum[0], datum[1]))
    elif "token_map" in datum:
      examples[-1].features[idx] = datum
    else:
      examples[-1].results[idx] = datum

  # Construct prediction objects.
  tf.compat.v1.logging.info("Computing predictions...")
  summary_dict = {}
  nq_pred_dict = {}
  for e in examples:
    summary = compute_predictions(e)
    summary_dict[e.example_id] = summary
    nq_pred_dict[e.example_id] = summary.predicted_label
    if len(nq_pred_dict) % 100 == 0:
      tf.compat.v1.logging.info("Examples processed: %d", len(nq_pred_dict))
  tf.compat.v1.logging.info("Done computing predictions.")

  return nq_pred_dict
//  [python]
def _build_train(self):
    print("Build train graph")
    logits = self._model(self.x_train, True)
    log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=self.y_train)
    self.loss = tf.reduce_mean(log_probs)

    self.train_preds = tf.argmax(logits, axis=1)
    self.train_preds = tf.to_int32(self.train_preds)
    self.train_acc = tf.equal(self.train_preds, self.y_train)
    self.train_acc = tf.to_int32(self.train_acc)
    self.train_acc = tf.reduce_sum(self.train_acc)

    tf_variables = [var
        for var in tf.trainable_variables() if var.name.startswith(self.name)]
    self.num_vars = count_model_params(tf_variables)
    print("-" * 80)
    for var in tf_variables:
      print(var)

    self.global_step = tf.Variable(
      0, dtype=tf.int32, trainable=False, name="global_step")
    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      self.loss,
      tf_variables,
      self.global_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)
//  [python]
def _build_train(self):
    print("Build train graph")
    logits = self._model(self.x_train, True)
    log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=self.y_train)
    self.loss = tf.reduce_mean(log_probs)

    self.train_preds = tf.argmax(logits, axis=1)
    self.train_preds = tf.cast(self.train_preds, tf.int32)
    self.train_acc = tf.equal(self.train_preds, self.y_train)
    self.train_acc = tf.cast(self.train_acc, tf.int32)
    self.train_acc = tf.reduce_sum(self.train_acc)

    tf_variables = [var
        for var in tf.trainable_variables() if var.name.startswith(self.name)]
    self.num_vars = count_model_params(tf_variables)
    print("-" * 80)
    for var in tf_variables:
      print(var)

    self.global_step = tf.Variable(
      0, dtype=tf.int32, trainable=False, name="global_step")
    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      self.loss,
      tf_variables,
      self.global_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
    #with tf.device("/device:GPU:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])

        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(_pre_process, x_valid_shuffle,
                                    back_prop=False)

    logits = self._model(x_valid_shuffle, False, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.to_int32(valid_shuffle_preds)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.to_int32(self.valid_shuffle_acc)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)