jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def add_image_summaries(features, params):
    """Add sample summaries for train frames and flows."""

    # Set up frames and flows if available
    if params.use_flows and params.use_frames:
        frames = tf.cast(features[0], tf.float32)
        flows = tf.cast(features[1], tf.float32)
    elif params.use_frames:
        frames = tf.cast(features, tf.float32); flows = None
    elif params.use_flows:
        flows = tf.cast(features, tf.float32); frames = None

    def _flow_to_image(flow):
        """Convert flow representation for visualisation"""
        def _cart_to_polar(x, y):
            """Conversion from cart to polar coordinates."""
            theta = tf.atan2(y, x)
            rho = tf.sqrt(tf.add(tf.pow(x, 2), tf.pow(y, 2)))
            return theta, rho
        def _normalize(x, min, max):
            """Minmax normalization."""
            norm = tf.divide(
                tf.subtract(x, tf.reduce_min(x)),
                tf.subtract(tf.reduce_max(x), tf.reduce_min(x)))
            norm = tf.add(tf.multiply(norm, max-min), min)
            return norm
        ang, mag = _cart_to_polar(flow[:,:,0], flow[:,:,1])
        hsv_0 = tf.multiply(ang, tf.divide(180.0, tf.divide(math.pi, 2.0)))
        hsv_1 = tf.fill(tf.shape(hsv_0), 255.0)
        hsv_2 = _normalize(mag, 0, 255)
        res = tf.stack([hsv_0, hsv_1, hsv_2], axis=2)/255.0
        res = tf.image.hsv_to_rgb(res)
        return res

    if params.use_sequence_input:
        if frames is not None:
            for i in range(5):
                tf.compat.v1.summary.image(name="frame" + str(i),
                                tensor=frames[i],
                                max_outputs=params.sequence_length)
        if flows is not None:
            for i in range(5):
                temp = tf.map_fn(
                    fn=_flow_to_image,
                    elems=flows[i,0:params.sequence_length])
                tf.compat.v1.summary.image(name="flow" + str(i),
                                tensor=temp,
                                max_outputs=params.sequence_length)
    else:
        if frames is not None:
            tf.compat.v1.summary.image('frames', frames, max_outputs=10)
        if flows is not None:
            flows = tf.map_fn(_flow_to_image, flows[0:9])
            tf.compat.v1.summary.image('flows', flows, max_outputs=10)
//  [python]
def sample_video_frames(
    batch_video_matrix: tf.Tensor,
    num_frames: tf.Tensor,
    random_frames: bool = True,
    num_sample_frames: int = 25,
):
  """Preprocesses input to sample frames."""

  # Sample random frames / random sequence.
  num_frames = t.cast(num_frames, tf.float32)
  if random_frames:
    batch_video_matrix = sample_random_frames(
        batch_video_matrix, num_frames, num_sample_frames
    )
  else:
    batch_video_matrix = sample_random_sequence(
        batch_video_matrix, num_frames, num_sample_frames
    )
  return batch_video_matrix
//  [python]
def sample_video_frames(
    batch_video_matrix: tf.Tensor,
    num_frames: tf.Tensor,
    random_frames: bool = True,
    num_sample_frames: int = 25,
):
  """Preprocesses input to sample frames."""

  # Sample random frames / random sequence.
  num_frames = tf.cast(num_frames, tf.float32)
  if random_frames:
    batch_video_matrix = sample_random_frames(
        batch_video_matrix, num_frames, num_sample_frames
    )
  else:
    batch_video_matrix = sample_random_sequence(
        batch_video_matrix, num_frames, num_sample_frames
    )
  return batch_video_matrix
//  [python]
def fit_task_critics(self, mb_states, mb_actions,
                       mb_next_states, mb_next_actions,
                       mb_rewards, mb_discounts,
                       level_ids):
    """Updates per-level critic parameters.

    Args:
      mb_states: Batch of states.
      mb_actions: Batch of actions.
      mb_next_states: Batch of next states.
      mb_next_actions: Batch of next actions from training policy.
      mb_rewards: Batch of rewards.
      mb_discounts: Batch of masks indicating the end of the episodes.
      level_ids: Batch of level ids

    Returns:
      Dictionary with information to track.
    """
    mb_rewards = self.reward_normalizer.normalize_target(mb_rewards)
    trainable_variables = (self.encoder_per_level.trainable_variables
                           + self.task_critic_one.trainable_variables)

    next_action_indices = tf.stack([
        tf.range(tf.shape(mb_next_actions)[0],
                 dtype=tf.int32), level_ids * self.output_dim_level +
        tf.cast(mb_next_actions, dtype=tf.int32)
    ],
                                   axis=-1)

    action_indices = tf.stack([
        tf.range(tf.shape(mb_actions)[0], dtype=tf.int32),
        level_ids * self.output_dim_level + tf.cast(mb_actions, dtype=tf.int32)
    ],
                              axis=-1)
    level_ids = tf.stack([
        tf.range(tf.shape(mb_next_actions)[0],
                 dtype=tf.int32), tf.cast(level_ids, dtype=tf.int32)
    ],
                         axis=-1)

    if 'parallelPerLevel' in self.rep_learn_keywords:
      next_states = [self.encoder_per_level_target(mb_next_states[0])]
      next_q1, next_q2 = self.task_critic_target_one(
          next_states[0], actions=None)
      target_q = tf.expand_dims(
          mb_rewards, 1) + self.discount * tf.expand_dims(
              mb_discounts, 1) * tf.minimum(next_q1, next_q2)

    target_q = tf.gather_nd(target_q, indices=next_action_indices)

    with tf.GradientTape(watch_accessed_variables=False) as tape:
      tape.watch(trainable_variables)

      states = [self.encoder_per_level(mb_states[0])]
      q1_all, q2_all = self.task_critic_one(states[0], actions=None)

      q = tf.minimum(q1_all, q2_all)

      critic_loss = (
          tf.losses.mean_squared_error(
              target_q, tf.gather_nd(q1_all, indices=action_indices)) +
          tf.losses.mean_squared_error(
              target_q, tf.gather_nd(q2_all, indices=action_indices)))

    critic_grads = tape.gradient(critic_loss, trainable_variables)

    self.task_critic_optimizer.apply_gradients(zip(critic_grads,
                                                   trainable_variables))

    critic.soft_update(
        self.encoder_per_level, self.encoder_per_level_target, tau=self.tau)
    critic.soft_update(
        self.task_critic_one, self.task_critic_target_one, tau=self.tau)

    gn = tf.reduce_mean(
        [tf.linalg.norm(v) for v in critic_grads if v is not None])

    return {
        'avg_level_critic_loss': tf.reduce_mean(critic_loss),
        'avg_q': tf.reduce_mean(q),
        'level_critic_grad_norm': gn
    }
//  [python]
def call(self, inputs):
    layer_output = inputs
    for i in range(self._num_blocks):
      layer_input = layer_output
      intermediate_output = self._intermediate_dense[i](layer_input)
      intermediate_output = self._inner_activation_layers[i](
          intermediate_output)
      if self._use_gate:
        gated_linear = self._gate_dense[i](layer_input)
        intermediate_output = intermediate_output * gated_linear

      layer_output = self._output_dense[i](intermediate_output)
      if self._dropout_position == "before_residual":
        layer_output = self._output_dropout[i](layer_output)

      # During mixed precision training, `layer_input` may be from layer norm.
      # If so, it is always fp32. Cast layer_output to fp32 for the subsequent
      # add.
      if layer_input.dtype == tf.float32:
        layer_output = tf.to_float(layer_output)
      if self._apply_output_layer_norm:
        layer_output = self._output_layer_norm[i](layer_output + layer_input)
      if self._dropout_position == "after_residual":
        layer_output = self._output_dropout[i](layer_output)

    return layer_output
//  [python]
def call(self, inputs):
    layer_output = inputs
    for i in range(self._num_blocks):
      layer_input = layer_output
      intermediate_output = self._intermediate_dense[i](layer_input)
      intermediate_output = self._inner_activation_layers[i](
          intermediate_output)
      if self._use_gate:
        gated_linear = self._gate_dense[i](layer_input)
        intermediate_output = intermediate_output * gated_linear

      layer_output = self._output_dense[i](intermediate_output)
      if self._dropout_position == "before_residual":
        layer_output = self._output_dropout[i](layer_output)

      # During mixed precision training, `layer_input` may be from layer norm.
      # If so, it is always fp32. Cast layer_output to fp32 for the subsequent
      # add.
      if layer_input.dtype == tf.float32:
        layer_output = tf.cast(layer_output, tf.float32)
      if self._apply_output_layer_norm:
        layer_output = self._output_layer_norm[i](layer_output + layer_input)
      if self._dropout_position == "after_residual":
        layer_output = self._output_dropout[i](layer_output)

    return layer_output
//  [python]
logits=[]
costs=[]
i=0
accuracy=0
# for output in outputs:
for i in range(0, max_word_length):
	output=outputs[-i-1]
	weights = tf.Variable(tf.random_uniform([num_hidden, classes], minval=-1. / width, maxval=1./width), name="weights_%d"%i)
	bias = tf.Variable(tf.random_uniform([classes], minval=-1. / width, maxval=1. / width), name="bias_dense_%d"%i)
	y_ = outputY = tf.matmul(output, weights, name="dense_%d" % i) + bias

	cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_, labels=y[:,i,:]), name="cost")  # prediction, target
	costs.append(cost)
	logits.append(y_)

	correct_pred = tf.equal(tf.argmax(outputY, 1), tf.argmax(y[:,i], 1))
	accuraci = tf.reduce_mean(tf.to_float(correct_pred))
	accuracy+= accuraci
//  [python]
logits=[]
costs=[]
i=0
accuracy=0
# for output in outputs:
for i in range(0, max_word_length):
	output=outputs[-i-1]
	weights = tf.Variable(tf.random_uniform([num_hidden, classes], minval=-1. / width, maxval=1./width), name="weights_%d"%i)
	bias = tf.Variable(tf.random_uniform([classes], minval=-1. / width, maxval=1. / width), name="bias_dense_%d"%i)
	y_ = outputY = tf.matmul(output, weights, name="dense_%d" % i) + bias

	cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_, labels=y[:,i,:]), name="cost")  # prediction, target
	costs.append(cost)
	logits.append(y_)

	correct_pred = tf.equal(tf.argmax(outputY, 1), tf.argmax(y[:,i], 1))
	accuraci = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
	accuracy+= accuraci
//  [python]
def fpn_map_rois_to_levels(boxes):
    """
    Assign boxes to level 2~5.

    Args:
        boxes (nx4):

    Returns:
        [tf.Tensor]: 4 tensors for level 2-5. Each tensor is a vector of indices of boxes in its level.
        [tf.Tensor]: 4 tensors, the gathered boxes in each level.

    Be careful that the returned tensor could be empty.
    """
    sqrtarea = tf.sqrt(tf_area(boxes))
    level = tf.floor(
        4 + tf.math.log(sqrtarea * (1. / 224) + 1e-6) * (1.0 / np.log(2)))

    # RoI levels range from 2~5 (not 6)
    level_ids = [
        tf.where(level <= 2),
        tf.where(tf.equal(level, 3)),   # == is not supported
        tf.where(tf.equal(level, 4)),
        tf.where(level >= 5)]
    level_ids = [tf.reshape(x, [-1], name='roi_level{}_id'.format(i + 2))
                 for i, x in enumerate(level_ids)]
    num_in_levels = [tf.size(x, name='num_roi_level{}'.format(i + 2))
                     for i, x in enumerate(level_ids)]
    add_moving_summary(*num_in_levels)

    level_boxes = [tf.gather(boxes, ids) for ids in level_ids]
    return level_ids, level_boxes
//  [python]
def fpn_map_rois_to_levels(boxes):
    """
    Assign boxes to level 2~5.

    Args:
        boxes (nx4):

    Returns:
        [tf.Tensor]: 4 tensors for level 2-5. Each tensor is a vector of indices of boxes in its level.
        [tf.Tensor]: 4 tensors, the gathered boxes in each level.

    Be careful that the returned tensor could be empty.
    """
    sqrtarea = tf.sqrt(tf_area(boxes))
    level = tf.cast(tf.floor(
        4 + tf.math.log(sqrtarea * (1. / 224) + 1e-6) * (1.0 / np.log(2))), tf.int32)

    # RoI levels range from 2~5 (not 6)
    level_ids = [
        tf.where(level <= 2),
        tf.where(tf.equal(level, 3)),   # == is not supported
        tf.where(tf.equal(level, 4)),
        tf.where(level >= 5)]
    level_ids = [tf.reshape(x, [-1], name='roi_level{}_id'.format(i + 2))
                 for i, x in enumerate(level_ids)]
    num_in_levels = [tf.size(x, name='num_roi_level{}'.format(i + 2))
                     for i, x in enumerate(level_ids)]
    add_moving_summary(*num_in_levels)

    level_boxes = [tf.gather(boxes, ids) for ids in level_ids]
    return level_ids, level_boxes
//  [python]
def accuracy(margin=0.0):
    """
    Compute the relative number of pairs with a score in the margin, ie. #{pairs | |y_true - y_pred| < m}
    """

    def _accuracy(y_true, y_pred):
        y_true = tf.matmul(y_true, y_true, transpose_b=True)
        return tf.reduce_mean(tf.to_float(tf.abs(y_true - y_pred) < margin))

    return _accuracy
//  [python]
def accuracy(margin=0.0):
    """
    Compute the relative number of pairs with a score in the margin, ie. #{pairs | |y_true - y_pred| < m}
    """

    def _accuracy(y_true, y_pred):
        y_true = tf.matmul(y_true, y_true, transpose_b=True)
        return tf.reduce_mean(tf.cast(tf.abs(y_true - y_pred) < margin, y_pred.dtype))

    return _accuracy
//  [python]
def _resize_and_pad_masks(self, mask, image_info):
    """Resizes masks to match the original image shape and pads to`output_size`.

    Args:
      mask: a padded mask tensor.
      image_info: a tensor that holds information about original and
        preprocessed images.
    Returns:
      resized and padded masks: tf.Tensor.
    """
    rescale_size = tf.to_int32(
        tf.math.ceil(image_info[1, :] / image_info[2, :]))
    image_shape = tf.to_int32(image_info[0, :])
    offsets = tf.to_int32(image_info[3, :])

    mask = tf.image.resize(
        mask,
        rescale_size,
        method='bilinear')
    mask = tf.image.crop_to_bounding_box(
        mask,
        offsets[0], offsets[1],
        image_shape[0],
        image_shape[1])
    mask = tf.image.pad_to_bounding_box(
        mask, 0, 0,
        self._config_dict['output_size'][0],
        self._config_dict['output_size'][1])
    return mask
//  [python]
def _resize_and_pad_masks(self, mask, image_info):
    """Resizes masks to match the original image shape and pads to`output_size`.

    Args:
      mask: a padded mask tensor.
      image_info: a tensor that holds information about original and
        preprocessed images.
    Returns:
      resized and padded masks: tf.Tensor.
    """
    rescale_size = tf.cast(
        tf.math.ceil(image_info[1, :] / image_info[2, :]), tf.int32)
    image_shape = tf.cast(image_info[0, :], tf.int32)
    offsets = tf.cast(image_info[3, :], tf.int32)

    mask = tf.image.resize(
        mask,
        rescale_size,
        method='bilinear')
    mask = tf.image.crop_to_bounding_box(
        mask,
        offsets[0], offsets[1],
        image_shape[0],
        image_shape[1])
    mask = tf.image.pad_to_bounding_box(
        mask, 0, 0,
        self._config_dict['output_size'][0],
        self._config_dict['output_size'][1])
    return mask
//  [python]
def _find_closest_center_per_pixel(centers: tf.Tensor,
                                   center_offsets: tf.Tensor) -> tf.Tensor:
  """Assigns all pixels to their closest center.

  Args:
    centers: A tf.Tensor of shape [N, 2] containing N centers with coordinate
      order (y, x).
    center_offsets: A tf.Tensor of shape [height, width, 2].
  Returns:
    A tf.Tensor of shape [height, width] containing the index of the closest
      center, per pixel.
  """
  height = tf.shape(center_offsets)[0]
  width = tf.shape(center_offsets)[1]

  x_coord, y_coord = tf.meshgrid(tf.range(width), tf.range(height))
  coord = tf.stack([y_coord, x_coord], axis=-1)

  center_per_pixel = t.cast(coord, tf.float32) + center_offsets

  # centers: [N, 2] -> [N, 1, 2].
  # center_per_pixel: [H, W, 2] -> [1, H*W, 2].
  centers = t.cast(tf.expand_dims(centers, 1), tf.float32)
  center_per_pixel = tf.reshape(center_per_pixel, [height*width, 2])
  center_per_pixel = tf.expand_dims(center_per_pixel, 0)

  # distances: [N, H*W].
  distances = tf.norm(centers - center_per_pixel, axis=-1)

  return tf.reshape(tf.argmin(distances, axis=0), [height, width])
//  [python]
def _find_closest_center_per_pixel(centers: tf.Tensor,
                                   center_offsets: tf.Tensor) -> tf.Tensor:
  """Assigns all pixels to their closest center.

  Args:
    centers: A tf.Tensor of shape [N, 2] containing N centers with coordinate
      order (y, x).
    center_offsets: A tf.Tensor of shape [height, width, 2].
  Returns:
    A tf.Tensor of shape [height, width] containing the index of the closest
      center, per pixel.
  """
  height = tf.shape(center_offsets)[0]
  width = tf.shape(center_offsets)[1]

  x_coord, y_coord = tf.meshgrid(tf.range(width), tf.range(height))
  coord = tf.stack([y_coord, x_coord], axis=-1)

  center_per_pixel = tf.cast(coord, tf.float32) + center_offsets

  # centers: [N, 2] -> [N, 1, 2].
  # center_per_pixel: [H, W, 2] -> [1, H*W, 2].
  centers = tf.cast(tf.expand_dims(centers, 1), tf.float32)
  center_per_pixel = tf.reshape(center_per_pixel, [height*width, 2])
  center_per_pixel = tf.expand_dims(center_per_pixel, 0)

  # distances: [N, H*W].
  distances = tf.norm(centers - center_per_pixel, axis=-1)

  return tf.reshape(tf.argmin(distances, axis=0), [height, width])
//  [python]
def _resize_and_pad_offset_mask(self, mask, image_info):
    """Rescales and resizes offset masks and pads to`output_size`.

    Args:
      mask: a padded offset mask tensor.
      image_info: a tensor that holds information about original and
        preprocessed images.
    Returns:
      rescaled, resized and padded masks: tf.Tensor.
    """
    rescale_size = tf.to_int32(
        tf.math.ceil(image_info[1, :] / image_info[2, :]))
    image_shape = tf.to_int32(image_info[0, :])
    offsets = tf.to_int32(image_info[3, :])

    mask = mask_ops.resize_and_rescale_offsets(
        tf.expand_dims(mask, axis=0),
        rescale_size)[0]
    mask = tf.image.crop_to_bounding_box(
        mask,
        offsets[0], offsets[1],
        image_shape[0],
        image_shape[1])
    mask = tf.image.pad_to_bounding_box(
        mask, 0, 0,
        self._config_dict['output_size'][0],
        self._config_dict['output_size'][1])
    return mask
//  [python]
def _resize_and_pad_offset_mask(self, mask, image_info):
    """Rescales and resizes offset masks and pads to`output_size`.

    Args:
      mask: a padded offset mask tensor.
      image_info: a tensor that holds information about original and
        preprocessed images.
    Returns:
      rescaled, resized and padded masks: tf.Tensor.
    """
    rescale_size = tf.cast(
        tf.math.ceil(image_info[1, :] / image_info[2, :]), tf.int32)
    image_shape = tf.cast(image_info[0, :], tf.int32)
    offsets = tf.cast(image_info[3, :], tf.int32)

    mask = mask_ops.resize_and_rescale_offsets(
        tf.expand_dims(mask, axis=0),
        rescale_size)[0]
    mask = tf.image.crop_to_bounding_box(
        mask,
        offsets[0], offsets[1],
        image_shape[0],
        image_shape[1])
    mask = tf.image.pad_to_bounding_box(
        mask, 0, 0,
        self._config_dict['output_size'][0],
        self._config_dict['output_size'][1])
    return mask
//  [python]
def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    src_len = shape_list(mask)[1]
    tgt_len = tgt_len if tgt_len is not None else src_len
    one_cst = tf.constant(1.0)
    mask = t.cast(mask, dtype=one_cst.dtype)
    expanded_mask = tf.tile(mask[:, None, None, :], (1, 1, tgt_len, 1))

    return (one_cst - expanded_mask) * LARGE_NEGATIVE
//  [python]
def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    src_len = shape_list(mask)[1]
    tgt_len = tgt_len if tgt_len is not None else src_len
    one_cst = tf.constant(1.0)
    mask = tf.cast(mask, dtype=one_cst.dtype)
    expanded_mask = tf.tile(mask[:, None, None, :], (1, 1, tgt_len, 1))

    return (one_cst - expanded_mask) * LARGE_NEGATIVE
//  [python]
def infer_length(seq, eos_ix, time_major=False, dtype=tf.int32):
    """
    compute length given output indices and eos code
    :param seq: tf matrix [time,batch] if time_major else [batch,time]
    :param eos_ix: integer index of end-of-sentence token
    :returns: lengths, int32 vector of shape [batch]
    """
    axis = 0 if time_major else 1
    is_eos = t.cast(t.equal(seq, eos_ix), dtype)
    count_eos = t.cumsum(is_eos, axis=axis, exclusive=True)
    lengths = t.reduce_sum(t.cast(t.equal(count_eos, 0), dtype), axis=axis)
    return lengths
//  [python]
def infer_length(seq, eos_ix, time_major=False, dtype=tf.int32):
    """
    compute length given output indices and eos code
    :param seq: tf matrix [time,batch] if time_major else [batch,time]
    :param eos_ix: integer index of end-of-sentence token
    :returns: lengths, int32 vector of shape [batch]
    """
    axis = 0 if time_major else 1
    is_eos = tf.cast(tf.equal(seq, eos_ix), dtype)
    count_eos = tf.cumsum(is_eos, axis=axis, exclusive=True)
    lengths = tf.reduce_sum(tf.cast(tf.equal(count_eos, 0), dtype), axis=axis)
    return lengths
//  [python]
def train_step(inputs):
    with tf.GradientTape() as tape:
        # pass inputs to SAM model
        outputs = sam(
            pixel_values=inputs["pixel_values"],
            input_boxes=inputs["input_boxes"],
            multimask_output=False,
            training=True,
        )

        predicted_masks = tf.squeeze(outputs.pred_masks, 1)
        ground_truth_masks = tf.cast(inputs["ground_truth_mask"], tf.float32)

        # calculate loss over predicted and ground truth masks
        loss = dice_loss(tf.expand_dims(ground_truth_masks, 1), predicted_masks)
        # update trainable variables
        trainable_vars = sam.trainable_variables
        grads = tape.gradient(loss, trainable_vars)
        optimizer.apply_gradients(zip(grads, trainable_vars))

        return loss
//  [python]
def masked_argmax(
    input_tensor: types.Tensor,
    mask: types.Tensor,
    output_type: tf.DType = tf.int32,
) -> types.Tensor:
  """Computes the argmax where the allowed elements are given by a mask.

  If a row of `mask` contains all zeros, then this method will return -1 for the
  corresponding row of `input_tensor`.

  Args:
    input_tensor: Rank-2 Tensor of floats.
    mask: 0-1 valued Tensor of the same shape as input.
    output_type: Integer type of the output.

  Returns:
    A Tensor of rank 1 and type `output_type`, with the masked argmax of every
    row of `input_tensor`.
  """
  input_tensor.shape.assert_is_compatible_with(mask.shape)
  neg_inf = tf.constant(-float('Inf'), input_tensor.dtype)
  modified_input = tf.compat.v2.where(
      tf.cast(mask, tf.bool), input_tensor, neg_inf
  )
  argmax_tensor = tf.argmax(modified_input, axis=-1, output_type=output_type)
  # Replace results for invalid mask rows with -1.
  reduce_mask = tf.cast(tf.reduce_max(mask, axis=1), tf.bool)
  neg_one = tf.constant(-1, output_type)
  return tf.compat.v2.where(reduce_mask, argmax_tensor, neg_one)
//  [python]
def rnn_model_fn(features, labels, mode):  # pylint: disable=unused-argument
  """Model function for a RNN."""

  # Define RNN architecture using tf.keras.layers.
  x = features['x']
  x = tf.reshape(x, [-1, SEQ_LEN])
  input_layer = x[:, :-1]
  input_one_hot = tf.one_hot(input_layer, 256)
  lstm = tf.keras.layers.LSTM(256, return_sequences=True)(input_one_hot)
  logits = tf.keras.layers.Dense(256)(lstm)

  # Calculate loss as a vector (to support microbatches in DP-SGD).
  vector_loss = tf.nn.softmax_cross_entropy_with_logits(
      labels=tf.one_hot(x[:, 1:], 256),
      logits=logits)
  # Define mean of loss across minibatch (for reporting through tf.Estimator).
  scalar_loss = tf.reduce_mean(vector_loss)

  # Configure the training op (for TRAIN mode).
  if mode == tf_estimator.ModeKeys.TRAIN:
    if FLAGS.dpsgd:

      optimizer = dp_optimizer.DPAdamGaussianOptimizer(
          l2_norm_clip=FLAGS.l2_norm_clip,
          noise_multiplier=FLAGS.noise_multiplier,
          num_microbatches=FLAGS.microbatches,
          learning_rate=FLAGS.learning_rate,
          unroll_microbatches=True)
      opt_loss = vector_loss
    else:
      optimizer = tf.compat.v1.train.AdamOptimizer(
          learning_rate=FLAGS.learning_rate)
      opt_loss = scalar_loss
    global_step = tf.compat.v1.train.get_global_step()
    train_op = optimizer.minimize(loss=opt_loss, global_step=global_step)
    return tf_estimator.EstimatorSpec(
        mode=mode, loss=scalar_loss, train_op=train_op)

  # Add evaluation metrics (for EVAL mode).
  elif mode == tf_estimator.ModeKeys.EVAL:
    eval_metric_ops = {
        'accuracy':
            tf.metrics.accuracy(
                labels=tf.cast(x[:, 1:], dtype=tf.int32),
                predictions=tf.argmax(input=logits, axis=2))
    }
    return tf_estimator.EstimatorSpec(
        mode=mode, loss=scalar_loss, eval_metric_ops=eval_metric_ops)
//  [python]
def rnn_model_fn(features, labels, mode):  # pylint: disable=unused-argument
  """Model function for a RNN."""

  # Define RNN architecture using tf.keras.layers.
  x = features['x']
  x = tf.reshape(x, [-1, SEQ_LEN])
  input_layer = x[:, :-1]
  input_one_hot = tf.one_hot(input_layer, 256)
  lstm = tf.keras.layers.LSTM(256, return_sequences=True)(input_one_hot)
  logits = tf.keras.layers.Dense(256)(lstm)

  # Calculate loss as a vector (to support microbatches in DP-SGD).
  vector_loss = tf.nn.softmax_cross_entropy_with_logits(
      labels=tf.cast(tf.one_hot(x[:, 1:], 256), dtype=tf.float32),
      logits=logits)
  # Define mean of loss across minibatch (for reporting through tf.Estimator).
  scalar_loss = tf.reduce_mean(vector_loss)

  # Configure the training op (for TRAIN mode).
  if mode == tf_estimator.ModeKeys.TRAIN:
    if FLAGS.dpsgd:

      optimizer = dp_optimizer.DPAdamGaussianOptimizer(
          l2_norm_clip=FLAGS.l2_norm_clip,
          noise_multiplier=FLAGS.noise_multiplier,
          num_microbatches=FLAGS.microbatches,
          learning_rate=FLAGS.learning_rate,
          unroll_microbatches=True)
      opt_loss = vector_loss
    else:
      optimizer = tf.compat.v1.train.AdamOptimizer(
          learning_rate=FLAGS.learning_rate)
      opt_loss = scalar_loss
    global_step = tf.compat.v1.train.get_global_step()
    train_op = optimizer.minimize(loss=opt_loss, global_step=global_step)
    return tf_estimator.EstimatorSpec(
        mode=mode, loss=scalar_loss, train_op=train_op)

  # Add evaluation metrics (for EVAL mode).
  elif mode == tf_estimator.ModeKeys.EVAL:
    eval_metric_ops = {
        'accuracy':
            tf.metrics.accuracy(
                labels=tf.cast(x[:, 1:], dtype=tf.int32),
                predictions=tf.argmax(input=logits, axis=2))
    }
    return tf_estimator.EstimatorSpec(
        mode=mode, loss=scalar_loss, eval_metric_ops=eval_metric_ops)
//  [python]
def _create_truncated_avg_masks(input_mask: tf.Tensor,
                                pool_strides: Sequence[int],
                                transforms: Sequence[tf.Tensor]):
  """Computes attention masks.

  For [1,1,1,0,0]

  Args:
    input_mask: Tensor of shape [batch_size, seq_length].
    pool_strides: Sequence of pooling strides for each layer.
    transforms: Sequence of off-diagonal matrices filling with 0.0 and
      1/pool_stride.

  Returns:
    attention_masks: Sequence of attention masks for each layer.
  """

  def create_2d_mask(from_length, mask):
    return tf.einsum('F,BT->BFT', tf.ones([from_length], dtype=mask.dtype),
                     mask)

  attention_masks = []
  seq_length = tf.shape(input_mask)[-1]
  layer_mask = t.cast(input_mask, dtype=_get_policy_dtype())
  for pool_stride, transform in zip(pool_strides, transforms):
    if pool_stride == 1:
      attention_masks.append(create_2d_mask(seq_length, layer_mask))
    else:
      pooled_seq_length = t.cast(
          t.cast(seq_length, tf.float32) / t.cast(pool_stride, tf.float32),
          tf.int32,
      )
      attention_masks.append(create_2d_mask(pooled_seq_length, layer_mask))

      layer_mask = t.cast(
          tf.einsum('BF,FT->BT', layer_mask, transform) > 0.0,
          dtype=layer_mask.dtype,
      )
      seq_length = pooled_seq_length
  del seq_length

  return attention_masks
//  [python]
def _create_truncated_avg_masks(input_mask: tf.Tensor,
                                pool_strides: Sequence[int],
                                transforms: Sequence[tf.Tensor]):
  """Computes attention masks.

  For [1,1,1,0,0]

  Args:
    input_mask: Tensor of shape [batch_size, seq_length].
    pool_strides: Sequence of pooling strides for each layer.
    transforms: Sequence of off-diagonal matrices filling with 0.0 and
      1/pool_stride.

  Returns:
    attention_masks: Sequence of attention masks for each layer.
  """

  def create_2d_mask(from_length, mask):
    return tf.einsum('F,BT->BFT', tf.ones([from_length], dtype=mask.dtype),
                     mask)

  attention_masks = []
  seq_length = tf.shape(input_mask)[-1]
  layer_mask = tf.cast(input_mask, dtype=_get_policy_dtype())
  for pool_stride, transform in zip(pool_strides, transforms):
    if pool_stride == 1:
      attention_masks.append(create_2d_mask(seq_length, layer_mask))
    else:
      pooled_seq_length = tf.cast(
          tf.cast(seq_length, tf.float32) / tf.cast(pool_stride, tf.float32),
          tf.int32,
      )
      attention_masks.append(create_2d_mask(pooled_seq_length, layer_mask))

      layer_mask = tf.cast(
          tf.einsum('BF,FT->BT', layer_mask, transform) > 0.0,
          dtype=layer_mask.dtype,
      )
      seq_length = pooled_seq_length
  del seq_length

  return attention_masks
//  [python]
def testDistribution(self, seed):
    segment_ids = [0, 0, 1, 1, 1, 2, 2, 2, 2]
    distr = [0] * len(segment_ids)
    for _ in range(1_000):
      result = utils.segment_random_index_shuffle(
          segment_ids=tf.convert_to_tensor(segment_ids, tf.int32), seed=seed)
      distr[result[3].numpy()] += 1
    distr = t.cast(distr, tf.float32) / 1_000.0
    std = math.sqrt(0.333 * 0.667 / 1_000.0)
    self.assertAllEqual(distr[:2], [0., 0.])
    self.assertAllClose(distr[2:5], [0.333] * 3, atol=5. * std, rtol=0.0)
    self.assertAllEqual(distr[5:], [0., 0., 0., 0.])
//  [python]
def testDistribution(self, seed):
    segment_ids = [0, 0, 1, 1, 1, 2, 2, 2, 2]
    distr = [0] * len(segment_ids)
    for _ in range(1_000):
      result = utils.segment_random_index_shuffle(
          segment_ids=tf.convert_to_tensor(segment_ids, tf.int32), seed=seed)
      distr[result[3].numpy()] += 1
    distr = tf.cast(distr, tf.float32) / 1_000.0
    std = math.sqrt(0.333 * 0.667 / 1_000.0)
    self.assertAllEqual(distr[:2], [0., 0.])
    self.assertAllClose(distr[2:5], [0.333] * 3, atol=5. * std, rtol=0.0)
    self.assertAllEqual(distr[5:], [0., 0., 0., 0.])
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor = None,
        causal_attention_mask: tf.Tensor = None,
        output_attentions: bool = None,
        encoder_hidden_states: tf.Tensor = None,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        """Input shape: Batch x Time x Channel"""

        batch_size = shape_list(hidden_states)[0]
        is_cross_attention = encoder_hidden_states is not None

        mixed_query_layer = self.q_proj(inputs=hidden_states)
        if is_cross_attention:
            mixed_key_layer = self.k_proj(inputs=encoder_hidden_states)
            mixed_value_layer = self.v_proj(inputs=encoder_hidden_states)
        else:
            mixed_key_layer = self.k_proj(inputs=hidden_states)
            mixed_value_layer = self.v_proj(inputs=hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = t.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        # apply the causal_attention_mask first
        if causal_attention_mask is not None:
            # Apply the causal attention mask (precomputed for all layers in TFCLIPModel call() function)
            attention_scores = tf.add(attention_scores, causal_attention_mask)

        if attention_mask is not None:
            # Apply the attention mask (precomputed for all layers in TFCLIPModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        _attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=_attention_probs)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, embed_dim)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.embed_dim))

        attention_output = self.out_proj(attention_output)
        # In TFBert, attention weights are returned after dropout.
        # However, in CLIP, they are returned before dropout.
        outputs = (attention_output, _attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor = None,
        causal_attention_mask: tf.Tensor = None,
        output_attentions: bool = None,
        encoder_hidden_states: tf.Tensor = None,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        """Input shape: Batch x Time x Channel"""

        batch_size = shape_list(hidden_states)[0]
        is_cross_attention = encoder_hidden_states is not None

        mixed_query_layer = self.q_proj(inputs=hidden_states)
        if is_cross_attention:
            mixed_key_layer = self.k_proj(inputs=encoder_hidden_states)
            mixed_value_layer = self.v_proj(inputs=encoder_hidden_states)
        else:
            mixed_key_layer = self.k_proj(inputs=hidden_states)
            mixed_value_layer = self.v_proj(inputs=hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        # apply the causal_attention_mask first
        if causal_attention_mask is not None:
            # Apply the causal attention mask (precomputed for all layers in TFCLIPModel call() function)
            attention_scores = tf.add(attention_scores, causal_attention_mask)

        if attention_mask is not None:
            # Apply the attention mask (precomputed for all layers in TFCLIPModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        _attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=_attention_probs)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, embed_dim)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.embed_dim))

        attention_output = self.out_proj(attention_output)
        # In TFBert, attention weights are returned after dropout.
        # However, in CLIP, they are returned before dropout.
        outputs = (attention_output, _attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def load_network(model_path):
    sess = tf.Session()
    images_pl = tf.placeholder(tf.float32, shape=[None, 160, 160, 3], name='input_image')
    images_norm = tf.map_fn(lambda frame: tf.image.per_image_standardization(frame), images_pl)
    train_mode = tf.placeholder(tf.bool)
    age_logits, gender_logits, _ = inception_resnet_v1.inference(images_norm, keep_probability=0.8,
                                                                 phase_train=train_mode,
                                                                 weight_decay=1e-5)
    gender = tf.argmax(tf.nn.softmax(gender_logits), 1)
    age_ = tf.constant([i for i in range(0, 101)], dtype=tf.float32)
    age = tf.reduce_sum(tf.multiply(tf.nn.softmax(age_logits), age_), axis=1)
    init_op = tf.group(tf.global_variables_initializer(),
                       tf.local_variables_initializer())
    sess.run(init_op)
    saver = tf.train.Saver()
    ckpt = tf.train.get_checkpoint_state(model_path)
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)
        print("restore model!")
    else:
        pass
    return sess,age,gender,train_mode,images_pl
//  [python]
def load_network(model_path):
    sess = tf.Session()
    images_pl = tf.placeholder(tf.float32, shape=[None, 160, 160, 3], name='input_image')
    images_norm = tf.map_fn(lambda frame: tf.image.per_image_standardization(frame), images_pl)
    train_mode = tf.placeholder(tf.bool)
    age_logits, gender_logits, _ = inception_resnet_v1.inference(images_norm, keep_probability=0.8,
                                                                 phase_train=train_mode,
                                                                 weight_decay=1e-5)
    gender = tf.argmax(tf.nn.softmax(gender_logits), 1)
    age_ = tf.cast(tf.constant([i for i in range(0, 101)]), tf.float32)
    age = tf.reduce_sum(tf.multiply(tf.nn.softmax(age_logits), age_), axis=1)
    init_op = tf.group(tf.global_variables_initializer(),
                       tf.local_variables_initializer())
    sess.run(init_op)
    saver = tf.train.Saver()
    ckpt = tf.train.get_checkpoint_state(model_path)
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)
        print("restore model!")
    else:
        pass
    return sess,age,gender,train_mode,images_pl
//  [python]
def cast_tensor_dict(dict_tensor, dtype):
  """Cast a nested dictionary of tensors to a data type.

  Args:
    dict_tensor: Input nested dictionary of tensors to cast type.
    dtype: The data type to cast to.

  Returns:
    A casted nesteddictionary of tensors.
  """
  casted_tensors = {}
  for k, v in traverse_util.flatten_dict(dict_tensor).items():
    casted_tensors[k] = t.cast(v, dtype)
  return traverse_util.unflatten_dict(casted_tensors)
//  [python]
def cast_tensor_dict(dict_tensor, dtype):
  """Cast a nested dictionary of tensors to a data type.

  Args:
    dict_tensor: Input nested dictionary of tensors to cast type.
    dtype: The data type to cast to.

  Returns:
    A casted nesteddictionary of tensors.
  """
  casted_tensors = {}
  for k, v in traverse_util.flatten_dict(dict_tensor).items():
    casted_tensors[k] = tf.cast(v, dtype)
  return traverse_util.unflatten_dict(casted_tensors)
//  [python]
def testPredictedRewards(self, batch_size, exploration_strategy, alpha):
    a_inv = []
    theta = []
    for k in range(self._num_actions):
      a_inv.append(np.linalg.inv(self._a_numpy[k] + np.eye(self._obs_dim)))
      theta.append(
          np.matmul(a_inv[k], self._b_numpy[k].reshape([self._obs_dim, 1]))
      )
    policy = linear_policy.LinearBanditPolicy(
        self._action_spec,
        self._a,
        self._b,
        self._num_samples_per_arm,
        self._time_step_spec,
        exploration_strategy,
        alpha=alpha,
        theta=None
        if alpha > 0.0
        else tf.squeeze(
            tf.cast(tf.stack(theta), dtype=self._obs_spec.dtype), axis=-1
        ),
        emit_policy_info=(policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,),
    )

    action_step = policy.action(self._time_step_batch(batch_size=batch_size))
    self.assertEqual(action_step.action.shape.as_list(), [batch_size])
    self.assertEqual(action_step.action.dtype, tf.int32)

    observation_numpy = np.array(
        range(batch_size * self._obs_dim), dtype=np.float32
    ).reshape([batch_size, self._obs_dim])

    predicted_rewards_expected = []
    for k in range(self._num_actions):
      est_mean_reward = np.matmul(observation_numpy, theta[k])
      predicted_rewards_expected.append(est_mean_reward)

    predicted_rewards_expected_array = np.stack(
        predicted_rewards_expected, axis=-1
    ).reshape(batch_size, self._num_actions)
    p_info = self.evaluate(action_step.info)
    self.assertAllClose(
        p_info.predicted_rewards_mean, predicted_rewards_expected_array
    )
//  [python]
def _generate_boxes(self):
    """Generates multiscale anchor boxes.

    Returns:
      a Tensor of shape [N, 4], representing anchor boxes of all levels
      concatenated together.
    """
    boxes_all = []
    for level in range(self.min_level, self.max_level + 1):
      boxes_l = []
      for scale in range(self.num_scales):
        for aspect_ratio in self.aspect_ratios:
          stride = 2 ** level
          intermidate_scale = 2 ** (scale / float(self.num_scales))
          base_anchor_size = self.anchor_size * stride * intermidate_scale
          aspect_x = aspect_ratio ** 0.5
          aspect_y = aspect_ratio ** -0.5
          half_anchor_size_x = base_anchor_size * aspect_x / 2.0
          half_anchor_size_y = base_anchor_size * aspect_y / 2.0
          x = tf.range(stride / 2, self.image_size[1], stride)
          y = tf.range(stride / 2, self.image_size[0], stride)
          xv, yv = tf.meshgrid(x, y)
          xv = tf.to_float(tf.reshape(xv, [-1]))
          yv = tf.to_float(tf.reshape(yv, [-1]))
          # Tensor shape Nx4.
          boxes = tf.stack([yv - half_anchor_size_y, xv - half_anchor_size_x,
                            yv + half_anchor_size_y, xv + half_anchor_size_x],
                           axis=1)
          boxes_l.append(boxes)
      # Concat anchors on the same level to tensor shape NxAx4.
      boxes_l = tf.stack(boxes_l, axis=1)
      boxes_l = tf.reshape(boxes_l, [-1, 4])
      boxes_all.append(boxes_l)
    return tf.concat(boxes_all, axis=0)
//  [python]
def _generate_boxes(self):
    """Generates multiscale anchor boxes.

    Returns:
      a Tensor of shape [N, 4], representing anchor boxes of all levels
      concatenated together.
    """
    boxes_all = []
    for level in range(self.min_level, self.max_level + 1):
      boxes_l = []
      for scale in range(self.num_scales):
        for aspect_ratio in self.aspect_ratios:
          stride = 2 ** level
          intermidate_scale = 2 ** (scale / float(self.num_scales))
          base_anchor_size = self.anchor_size * stride * intermidate_scale
          aspect_x = aspect_ratio ** 0.5
          aspect_y = aspect_ratio ** -0.5
          half_anchor_size_x = base_anchor_size * aspect_x / 2.0
          half_anchor_size_y = base_anchor_size * aspect_y / 2.0
          x = tf.range(stride / 2, self.image_size[1], stride)
          y = tf.range(stride / 2, self.image_size[0], stride)
          xv, yv = tf.meshgrid(x, y)
          xv = tf.cast(tf.reshape(xv, [-1]), dtype=tf.float32)
          yv = tf.cast(tf.reshape(yv, [-1]), dtype=tf.float32)
          # Tensor shape Nx4.
          boxes = tf.stack([yv - half_anchor_size_y, xv - half_anchor_size_x,
                            yv + half_anchor_size_y, xv + half_anchor_size_x],
                           axis=1)
          boxes_l.append(boxes)
      # Concat anchors on the same level to tensor shape NxAx4.
      boxes_l = tf.stack(boxes_l, axis=1)
      boxes_l = tf.reshape(boxes_l, [-1, 4])
      boxes_all.append(boxes_l)
    return tf.concat(boxes_all, axis=0)
//  [python]
def unpack_labels(self, labels, is_box=False):
    """Unpacks an array of labels into multiscales labels.

    Args:
      labels: labels to unpack.
      is_box: to unpack anchor boxes or not. If it is true, will unpack to 2D,
        otherwise, will unpack to 3D.

    Returns:
      unpacked_labels: a dictionary contains unpack labels in different levels.
    """
    unpacked_labels = collections.OrderedDict()
    count = 0
    for level in range(self.min_level, self.max_level + 1):
      feat_size_y = t.cast(self.image_size[0] / 2 ** level, t.int32)
      feat_size_x = t.cast(self.image_size[1] / 2 ** level, t.int32)
      steps = feat_size_y * feat_size_x * self.anchors_per_location
      if is_box:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [-1, 4])
      else:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [feat_size_y, feat_size_x, -1])
      count += steps
    return unpacked_labels
//  [python]
def unpack_labels(self, labels, is_box=False):
    """Unpacks an array of labels into multiscales labels.

    Args:
      labels: labels to unpack.
      is_box: to unpack anchor boxes or not. If it is true, will unpack to 2D,
        otherwise, will unpack to 3D.

    Returns:
      unpacked_labels: a dictionary contains unpack labels in different levels.
    """
    unpacked_labels = collections.OrderedDict()
    count = 0
    for level in range(self.min_level, self.max_level + 1):
      feat_size_y = tf.cast(self.image_size[0] / 2 ** level, tf.int32)
      feat_size_x = tf.cast(self.image_size[1] / 2 ** level, tf.int32)
      steps = feat_size_y * feat_size_x * self.anchors_per_location
      if is_box:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [-1, 4])
      else:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [feat_size_y, feat_size_x, -1])
      count += steps
    return unpacked_labels