jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _make_variables(self):
    """Creates the variables representing the parameters of the distribution."""
    channels = self.batch_shape.num_elements()
    filters = (1,) + self.num_filters + (1,)
    scale = self.init_scale ** (1 / (len(self.num_filters) + 1))
    self._matrices = []
    self._biases = []
    self._factors = []

    for i in range(len(self.num_filters) + 1):

      def matrix_initializer(i=i):
        init = log_expm1(1 / scale / filters[i + 1])
        init = t.cast(init, dtype=self.dtype)
        init = tf.broadcast_to(init, (channels, filters[i + 1], filters[i]))
        return init

      matrix = tf.Variable(matrix_initializer, name="matrix_{}".format(i))
      self._matrices.append(matrix)

      def bias_initializer(i=i):
        return tf.random.uniform(
            (channels, filters[i + 1], 1), -.5, .5, dtype=self.dtype)

      bias = tf.Variable(bias_initializer, name="bias_{}".format(i))
      self._biases.append(bias)

      if i < len(self.num_filters):

        def factor_initializer(i=i):
          return tf.zeros((channels, filters[i + 1], 1), dtype=self.dtype)

        factor = tf.Variable(factor_initializer, name="factor_{}".format(i))
        self._factors.append(factor)
//  [python]
def _make_variables(self):
    """Creates the variables representing the parameters of the distribution."""
    channels = self.batch_shape.num_elements()
    filters = (1,) + self.num_filters + (1,)
    scale = self.init_scale ** (1 / (len(self.num_filters) + 1))
    self._matrices = []
    self._biases = []
    self._factors = []

    for i in range(len(self.num_filters) + 1):

      def matrix_initializer(i=i):
        init = log_expm1(1 / scale / filters[i + 1])
        init = tf.cast(init, dtype=self.dtype)
        init = tf.broadcast_to(init, (channels, filters[i + 1], filters[i]))
        return init

      matrix = tf.Variable(matrix_initializer, name="matrix_{}".format(i))
      self._matrices.append(matrix)

      def bias_initializer(i=i):
        return tf.random.uniform(
            (channels, filters[i + 1], 1), -.5, .5, dtype=self.dtype)

      bias = tf.Variable(bias_initializer, name="bias_{}".format(i))
      self._biases.append(bias)

      if i < len(self.num_filters):

        def factor_initializer(i=i):
          return tf.zeros((channels, filters[i + 1], 1), dtype=self.dtype)

        factor = tf.Variable(factor_initializer, name="factor_{}".format(i))
        self._factors.append(factor)
//  [python]
def get_beta_accumulators(opt, dtype):
  local_step = t.cast(opt.iterations + 1, dtype)
  beta_1_t = t.cast(opt._get_hyper("beta_1"), dtype)
  beta_1_power = tf.math.pow(beta_1_t, local_step)
  beta_2_t = t.cast(opt._get_hyper("beta_2"), dtype)
  beta_2_power = tf.math.pow(beta_2_t, local_step)
  return (beta_1_power, beta_2_power)
//  [python]
def get_beta_accumulators(opt, dtype):
  local_step = tf.cast(opt.iterations + 1, dtype)
  beta_1_t = tf.cast(opt._get_hyper("beta_1"), dtype)
  beta_1_power = tf.math.pow(beta_1_t, local_step)
  beta_2_t = tf.cast(opt._get_hyper("beta_2"), dtype)
  beta_2_power = tf.math.pow(beta_2_t, local_step)
  return (beta_1_power, beta_2_power)
//  [python]
def _get_loss_multiplier_for_tensor(self, tensor, losses_mask):
    loss_multiplier_shape = tf.stack([-1] + [1] * (len(tensor.shape) - 1))
    return tf.to_float(tf.reshape(losses_mask, loss_multiplier_shape))
//  [python]
def _get_loss_multiplier_for_tensor(self, tensor, losses_mask):
    loss_multiplier_shape = tf.stack([-1] + [1] * (len(tensor.shape) - 1))
    return tf.cast(tf.reshape(losses_mask, loss_multiplier_shape), tf.float32)
//  [python]
def _classification_loss(pred_labels, gt_labels, num_matched_boxes):
  """Computes the classification loss.

  Computes the classification loss with hard negative mining.
  Args:
    pred_labels: a dict from index to tensor of predicted class. The shape
      of the tensor is [batch_size, num_anchors, num_classes].
    gt_labels: a list of tensor that represents the classification groundtruth
      targets. The shape is [batch_size, num_anchors, 1].
    num_matched_boxes: the number of anchors that are matched to a groundtruth
      targets. This is used as the loss normalizater.
  Returns:
    box_loss: a float32 representing total box regression loss.
  """
  keys = sorted(pred_labels.keys())
  batch_size = gt_labels[0].shape[0]
  cross_entropy = []
  for i, k in enumerate(keys):
    gt_label = gt_labels[i]
    pred_label = tf.reshape(
        pred_labels[k],
        gt_label.get_shape().as_list() + [ssd_constants.NUM_CLASSES])
    cross_entropy.append(
        tf.reshape(
            _softmax_cross_entropy(pred_label, gt_label), [batch_size, -1]))

  # Put the rest of the loss computation on one device to avoid excessive
  # communication inside topk_mask with spatial partition
  with tf.device(tf.contrib.tpu.core(0)):
    cross_entropy = tf.concat(cross_entropy, 1)
    gt_label = tf.concat([tf.reshape(l, [batch_size, -1]) for l in gt_labels],
                         1)
    mask = tf.greater(gt_label, 0)
    float_mask = tf.to_float(mask)

    # Hard example mining
    neg_masked_cross_entropy = cross_entropy * (1 - float_mask)

    num_neg_boxes = tf.minimum(
        tf.to_int32(num_matched_boxes) * ssd_constants.NEGS_PER_POSITIVE,
        ssd_constants.NUM_SSD_BOXES)
    top_k_neg_mask = topk_mask.topk_mask(
        neg_masked_cross_entropy,
        tf.tile(num_neg_boxes[:, tf.newaxis], (1, ssd_constants.NUM_SSD_BOXES)))

    class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)

    # TODO(taylorrobie): Confirm that normalizing by the number of boxes matches
    # reference
    return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def _classification_loss(pred_labels, gt_labels, num_matched_boxes):
  """Computes the classification loss.

  Computes the classification loss with hard negative mining.
  Args:
    pred_labels: a dict from index to tensor of predicted class. The shape
      of the tensor is [batch_size, num_anchors, num_classes].
    gt_labels: a list of tensor that represents the classification groundtruth
      targets. The shape is [batch_size, num_anchors, 1].
    num_matched_boxes: the number of anchors that are matched to a groundtruth
      targets. This is used as the loss normalizater.
  Returns:
    box_loss: a float32 representing total box regression loss.
  """
  keys = sorted(pred_labels.keys())
  batch_size = gt_labels[0].shape[0]
  cross_entropy = []
  for i, k in enumerate(keys):
    gt_label = gt_labels[i]
    pred_label = tf.reshape(
        pred_labels[k],
        gt_label.get_shape().as_list() + [ssd_constants.NUM_CLASSES])
    cross_entropy.append(
        tf.reshape(
            _softmax_cross_entropy(pred_label, gt_label), [batch_size, -1]))

  # Put the rest of the loss computation on one device to avoid excessive
  # communication inside topk_mask with spatial partition
  with tf.device(tf.contrib.tpu.core(0)):
    cross_entropy = tf.concat(cross_entropy, 1)
    gt_label = tf.concat([tf.reshape(l, [batch_size, -1]) for l in gt_labels],
                         1)
    mask = tf.greater(gt_label, 0)
    float_mask = tf.cast(mask, tf.float32)

    # Hard example mining
    neg_masked_cross_entropy = cross_entropy * (1 - float_mask)

    num_neg_boxes = tf.minimum(
        tf.to_int32(num_matched_boxes) * ssd_constants.NEGS_PER_POSITIVE,
        ssd_constants.NUM_SSD_BOXES)
    top_k_neg_mask = topk_mask.topk_mask(
        neg_masked_cross_entropy,
        tf.tile(num_neg_boxes[:, tf.newaxis], (1, ssd_constants.NUM_SSD_BOXES)))

    class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)

    # TODO(taylorrobie): Confirm that normalizing by the number of boxes matches
    # reference
    return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def read_tensor_from_image_file(file_name,
                                input_height=299,
                                input_width=299,
                                input_mean=0,
                                input_std=255):
  input_name = "file_reader"
  output_name = "normalized"
  file_reader = tf.read_file(file_name, input_name)
  if file_name.endswith(".png"):
    image_reader = tf.image.decode_png(
        file_reader, channels=3, name="png_reader")
  elif file_name.endswith(".gif"):
    image_reader = tf.squeeze(
        tf.image.decode_gif(file_reader, name="gif_reader"))
  elif file_name.endswith(".bmp"):
    image_reader = tf.image.decode_bmp(file_reader, name="bmp_reader")
  else:
    image_reader = tf.image.decode_jpeg(
        file_reader, channels=3, name="jpeg_reader")
  float_caster = tf.to_float(image_reader)
  dims_expander = tf.expand_dims(float_caster, 0)
  resized = tf.image.resize_bilinear(dims_expander, [input_height, input_width])
  normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
  sess = tf.Session()
  result = sess.run(normalized)

  return result
//  [python]
def read_tensor_from_image_file(file_name,
                                input_height=299,
                                input_width=299,
                                input_mean=0,
                                input_std=255):
  input_name = "file_reader"
  output_name = "normalized"
  file_reader = tf.read_file(file_name, input_name)
  if file_name.endswith(".png"):
    image_reader = tf.image.decode_png(
        file_reader, channels=3, name="png_reader")
  elif file_name.endswith(".gif"):
    image_reader = tf.squeeze(
        tf.image.decode_gif(file_reader, name="gif_reader"))
  elif file_name.endswith(".bmp"):
    image_reader = tf.image.decode_bmp(file_reader, name="bmp_reader")
  else:
    image_reader = tf.image.decode_jpeg(
        file_reader, channels=3, name="jpeg_reader")
  float_caster = tf.cast(image_reader, tf.float32)
  dims_expander = tf.expand_dims(float_caster, 0)
  resized = tf.image.resize_bilinear(dims_expander, [input_height, input_width])
  normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
  sess = tf.Session()
  result = sess.run(normalized)

  return result
//  [python]
def _cal_pos_emb(
        self,
        dense_layer: keras.layers.Dense,
        position_ids: tf.Tensor,
        num_buckets: int,
        max_distance: int,
    ):
        rel_pos_matrix = tf.expand_dims(position_ids, axis=-2) - tf.expand_dims(position_ids, axis=-1)
        rel_pos = self.relative_position_bucket(rel_pos_matrix, num_buckets, max_distance)
        rel_pos_one_hot = tf.one_hot(rel_pos, depth=num_buckets, dtype=self.compute_dtype)
        embedding = dense_layer(rel_pos_one_hot)
        # batch_size, seq_length, seq_length, num_heads --> batch_size, num_heads, seq_length, seq_length
        embedding = tf.transpose(embedding, [0, 3, 1, 2])
        embedding = tf.cast(embedding, dtype=self.compute_dtype)
        return embedding
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
    """Compute loss function.

    Args:
      prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
      target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
      weights: a float tensor of shape, either [batch_size, num_anchors,
        num_classes] or [batch_size, num_anchors, 1]. If the shape is
        [batch_size, num_anchors, 1], all the classses are equally weighted.

    Returns:
      loss: a float tensor of shape [batch_size, num_anchors, num_classes]
        representing the value of the loss function.
    """
    if self._bootstrap_type == 'soft':
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
    else:
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.to_float(
              tf.sigmoid(prediction_tensor) > 0.5)
    per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
        labels=bootstrap_target_tensor, logits=prediction_tensor))
    return per_entry_cross_ent * weights
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
    """Compute loss function.

    Args:
      prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
      target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
      weights: a float tensor of shape, either [batch_size, num_anchors,
        num_classes] or [batch_size, num_anchors, 1]. If the shape is
        [batch_size, num_anchors, 1], all the classses are equally weighted.

    Returns:
      loss: a float tensor of shape [batch_size, num_anchors, num_classes]
        representing the value of the loss function.
    """
    if self._bootstrap_type == 'soft':
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
    else:
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.cast(
              tf.sigmoid(prediction_tensor) > 0.5, tf.float32)
    per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
        labels=bootstrap_target_tensor, logits=prediction_tensor))
    return per_entry_cross_ent * weights
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      outputs = tf.nest.map_structure(
          lambda x: tf.to_float(x), outputs)

      # Computes per-replica loss.
      loss, cls_loss, box_loss, model_loss = self.build_losses(
          outputs=outputs, labels=labels, aux_losses=model.losses
      )
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}

    all_losses = {
        'total_loss': loss,
        'cls_loss': cls_loss,
        'box_loss': box_loss,
        'model_loss': model_loss,
    }
    if metrics:
      for m in metrics:
        m.update_state(all_losses[m.name])
        logs.update({m.name: m.result()})

    return logs
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      outputs = tf.nest.map_structure(
          lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      loss, cls_loss, box_loss, model_loss = self.build_losses(
          outputs=outputs, labels=labels, aux_losses=model.losses
      )
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}

    all_losses = {
        'total_loss': loss,
        'cls_loss': cls_loss,
        'box_loss': box_loss,
        'model_loss': model_loss,
    }
    if metrics:
      for m in metrics:
        m.update_state(all_losses[m.name])
        logs.update({m.name: m.result()})

    return logs
//  [python]
def _randomly_sample_a_seed_point_by_voxelization(point_locations, voxel_size):
  """Randomly samples a seed point from a random voxel."""
  voxel_size = t.convert_to_tensor(voxel_size, dtype=t.float32)
  start_location = t.reduce_min(point_locations, axis=0)
  voxel_xyz_indices = t.cast(
      t.math.floordiv(point_locations - start_location, voxel_size),
      dtype=t.int32)
  _, pooled_single_number_indices = voxel_utils.compute_pooled_voxel_indices(
      voxel_xyz_indices=voxel_xyz_indices, pooling_size=(1, 1, 1))
  unique_voxel_indices, point_indices = t.unique(
      t.reshape(pooled_single_number_indices, [-1]))
  random_voxel_index = t.random.uniform(
      [], minval=0, maxval=t.shape(unique_voxel_indices)[0], dtype=t.int32)
  potential_seed_points = t.cast(
      t.reshape(t.where(t.equal(point_indices, random_voxel_index)), [-1]),
      dtype=t.int32)
  random_seed_point_index = t.random.uniform(
      [], minval=0, maxval=t.shape(potential_seed_points)[0], dtype=t.int32)
  return potential_seed_points[random_seed_point_index]
//  [python]
def _randomly_sample_a_seed_point_by_voxelization(point_locations, voxel_size):
  """Randomly samples a seed point from a random voxel."""
  voxel_size = tf.convert_to_tensor(voxel_size, dtype=tf.float32)
  start_location = tf.reduce_min(point_locations, axis=0)
  voxel_xyz_indices = tf.cast(
      tf.math.floordiv(point_locations - start_location, voxel_size),
      dtype=tf.int32)
  _, pooled_single_number_indices = voxel_utils.compute_pooled_voxel_indices(
      voxel_xyz_indices=voxel_xyz_indices, pooling_size=(1, 1, 1))
  unique_voxel_indices, point_indices = tf.unique(
      tf.reshape(pooled_single_number_indices, [-1]))
  random_voxel_index = tf.random.uniform(
      [], minval=0, maxval=tf.shape(unique_voxel_indices)[0], dtype=tf.int32)
  potential_seed_points = tf.cast(
      tf.reshape(tf.where(tf.equal(point_indices, random_voxel_index)), [-1]),
      dtype=tf.int32)
  random_seed_point_index = tf.random.uniform(
      [], minval=0, maxval=tf.shape(potential_seed_points)[0], dtype=tf.int32)
  return potential_seed_points[random_seed_point_index]
//  [python]
def add_point_offset_bins(inputs, voxel_grid_cell_size, num_bins_x, num_bins_y,
                          num_bins_z):
  """Adds points offsets features to the inputs.

  Args:
    inputs: A dictionary containing input tensors.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
    num_bins_x: Number of bins in x dimension.
    num_bins_y: Number of bins in y dimension.
    num_bins_z: Number of bins in z dimension.
  """
  add_point_offsets(inputs=inputs, voxel_grid_cell_size=voxel_grid_cell_size)
  num_bins_coef_float = tf.convert_to_tensor(
      [[num_bins_x, num_bins_y, num_bins_z]], dtype=tf.float32)
  num_bins_coef_int = tf.convert_to_tensor(
      [[num_bins_x, num_bins_y, num_bins_z]], dtype=tf.int32)
  scaled_offsets = (inputs[standard_fields.InputDataFields.point_offsets] +
                    0.5) * num_bins_coef_float
  scaled_offsets = tf.minimum(
      tf.maximum(tf.to_int32(scaled_offsets), 0), num_bins_coef_int)
  scaled_offsets = (
      scaled_offsets[:, 0] + scaled_offsets[:, 1] * num_bins_x +
      scaled_offsets[:, 2] * num_bins_x * num_bins_y)
  inputs[standard_fields.InputDataFields.point_offset_bins] = tf.one_hot(
      indices=scaled_offsets,
      depth=(num_bins_x * num_bins_y * num_bins_z),
      dtype=tf.float32)
//  [python]
def add_point_offset_bins(inputs, voxel_grid_cell_size, num_bins_x, num_bins_y,
                          num_bins_z):
  """Adds points offsets features to the inputs.

  Args:
    inputs: A dictionary containing input tensors.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
    num_bins_x: Number of bins in x dimension.
    num_bins_y: Number of bins in y dimension.
    num_bins_z: Number of bins in z dimension.
  """
  add_point_offsets(inputs=inputs, voxel_grid_cell_size=voxel_grid_cell_size)
  num_bins_coef_float = tf.convert_to_tensor(
      [[num_bins_x, num_bins_y, num_bins_z]], dtype=tf.float32)
  num_bins_coef_int = tf.convert_to_tensor(
      [[num_bins_x, num_bins_y, num_bins_z]], dtype=tf.int32)
  scaled_offsets = (inputs[standard_fields.InputDataFields.point_offsets] +
                    0.5) * num_bins_coef_float
  scaled_offsets = tf.minimum(
      tf.maximum(tf.cast(scaled_offsets, dtype=tf.int32), 0), num_bins_coef_int)
  scaled_offsets = (
      scaled_offsets[:, 0] + scaled_offsets[:, 1] * num_bins_x +
      scaled_offsets[:, 2] * num_bins_x * num_bins_y)
  inputs[standard_fields.InputDataFields.point_offset_bins] = tf.one_hot(
      indices=scaled_offsets,
      depth=(num_bins_x * num_bins_y * num_bins_z),
      dtype=tf.float32)
//  [python]
def voxelize_point_to_view_correspondences(inputs, view_indices_2d_inputs,
                                           voxels_pad_or_clip_size,
                                           voxel_grid_cell_size):
  """Voxelizes the point to image correspondences.

  Args:
    inputs: A dictionary containing input tensors.
    view_indices_2d_inputs: Dictionary containing the name of the 2d views.
    voxels_pad_or_clip_size: Number of target voxels to pad or clip to. If None,
      it will not perform the voxel padding.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
  """
  for key in sorted(view_indices_2d_inputs):
    voxel_correspondence, _, _, _, _ = (
        voxel_utils.pointcloud_to_sparse_voxel_grid(
            points=tf.expand_dims(
                inputs[standard_fields.InputDataFields.point_positions],
                axis=0),
            features=t.cast(view_indices_2d_inputs[key], dtype=tf.float32),
            num_valid_points=tf.expand_dims(
                inputs[standard_fields.InputDataFields.num_valid_points],
                axis=0),
            grid_cell_size=voxel_grid_cell_size,
            voxels_pad_or_clip_size=voxels_pad_or_clip_size,
            segment_func=tf.math.unsorted_segment_mean))
    inputs[('%s/voxel_indices_2d' % key)] = t.cast(
        voxel_correspondence, dtype=tf.int32)
//  [python]
def voxelize_point_to_view_correspondences(inputs, view_indices_2d_inputs,
                                           voxels_pad_or_clip_size,
                                           voxel_grid_cell_size):
  """Voxelizes the point to image correspondences.

  Args:
    inputs: A dictionary containing input tensors.
    view_indices_2d_inputs: Dictionary containing the name of the 2d views.
    voxels_pad_or_clip_size: Number of target voxels to pad or clip to. If None,
      it will not perform the voxel padding.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
  """
  for key in sorted(view_indices_2d_inputs):
    voxel_correspondence, _, _, _, _ = (
        voxel_utils.pointcloud_to_sparse_voxel_grid(
            points=tf.expand_dims(
                inputs[standard_fields.InputDataFields.point_positions],
                axis=0),
            features=tf.cast(view_indices_2d_inputs[key], dtype=tf.float32),
            num_valid_points=tf.expand_dims(
                inputs[standard_fields.InputDataFields.num_valid_points],
                axis=0),
            grid_cell_size=voxel_grid_cell_size,
            voxels_pad_or_clip_size=voxels_pad_or_clip_size,
            segment_func=tf.math.unsorted_segment_mean))
    inputs[('%s/voxel_indices_2d' % key)] = tf.cast(
        voxel_correspondence, dtype=tf.int32)
//  [python]
def get_extended_attention_mask(self, attention_mask, input_shape):
        if attention_mask is None:
            attention_mask = tf.fill(input_shape, 1)

        # We create a 3D attention mask from a 2D tensor mask.
        # Sizes are [batch_size, 1, 1, to_seq_length]
        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
        # this attention mask is more simple than the triangular masking of causal attention
        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
        extended_attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.

        extended_attention_mask = tf.to_float(extended_attention_mask)
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        return extended_attention_mask
//  [python]
def get_extended_attention_mask(self, attention_mask, input_shape):
        if attention_mask is None:
            attention_mask = tf.fill(input_shape, 1)

        # We create a 3D attention mask from a 2D tensor mask.
        # Sizes are [batch_size, 1, 1, to_seq_length]
        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
        # this attention mask is more simple than the triangular masking of causal attention
        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
        extended_attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.

        extended_attention_mask = tf.cast(extended_attention_mask, tf.float32)
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        return extended_attention_mask
//  [python]
def remove_pointcloud_noise(mesh_inputs, view_indices_2d_inputs,
                            voxel_grid_cell_size, voxel_density_threshold):
  """Removes the noise (points that are low density).

  Args:
    mesh_inputs: A dictionary containing point properties.
    view_indices_2d_inputs: A dictionary containing view indices 2d.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
    voxel_density_threshold: Voxel density threshold. Points that belong to a
      voxel with less density than this threshold are removed.
  """
  points_position = mesh_inputs[standard_fields.InputDataFields.point_positions]
  num_points = tf.shape(points_position)[0]
  points_density = tf.ones([num_points, 1], dtype=tf.float32)
  if standard_fields.InputDataFields.num_valid_points in mesh_inputs:
    num_valid_points = mesh_inputs[
        standard_fields.InputDataFields.num_valid_points]
  else:
    num_valid_points = num_points
  (voxels_density, _, _, segment_ids, _) = (
      voxel_utils.pointcloud_to_sparse_voxel_grid(
          points=tf.expand_dims(points_position, axis=0),
          features=tf.expand_dims(points_density, axis=0),
          num_valid_points=tf.expand_dims(num_valid_points, axis=0),
          grid_cell_size=voxel_grid_cell_size,
          voxels_pad_or_clip_size=None,
          segment_func=tf.math.unsorted_segment_sum))
  voxels_density = tf.to_int32(tf.squeeze(voxels_density, axis=0))
  segment_ids = tf.squeeze(segment_ids, axis=0)
  points_density = tf.gather(voxels_density, segment_ids)
  valid_points = tf.greater_equal(
      tf.reshape(points_density, [-1]), voxel_density_threshold)
  for key in standard_fields.get_input_point_fields():
    if key != standard_fields.InputDataFields.num_valid_points:
      if key in mesh_inputs:
        mesh_inputs[key] = tf.boolean_mask(mesh_inputs[key], valid_points)
  mesh_inputs[
      standard_fields.InputDataFields.num_valid_points] = tf.math.reduce_sum(
          tf.to_int32(valid_points))
  for key in sorted(view_indices_2d_inputs):
    view_indices_2d_inputs[key] = tf.transpose(
        tf.boolean_mask(
            tf.transpose(view_indices_2d_inputs[key], [1, 0, 2]), valid_points),
        [1, 0, 2])
//  [python]
def remove_pointcloud_noise(mesh_inputs, view_indices_2d_inputs,
                            voxel_grid_cell_size, voxel_density_threshold):
  """Removes the noise (points that are low density).

  Args:
    mesh_inputs: A dictionary containing point properties.
    view_indices_2d_inputs: A dictionary containing view indices 2d.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
    voxel_density_threshold: Voxel density threshold. Points that belong to a
      voxel with less density than this threshold are removed.
  """
  points_position = mesh_inputs[standard_fields.InputDataFields.point_positions]
  num_points = tf.shape(points_position)[0]
  points_density = tf.ones([num_points, 1], dtype=tf.float32)
  if standard_fields.InputDataFields.num_valid_points in mesh_inputs:
    num_valid_points = mesh_inputs[
        standard_fields.InputDataFields.num_valid_points]
  else:
    num_valid_points = num_points
  (voxels_density, _, _, segment_ids, _) = (
      voxel_utils.pointcloud_to_sparse_voxel_grid(
          points=tf.expand_dims(points_position, axis=0),
          features=tf.expand_dims(points_density, axis=0),
          num_valid_points=tf.expand_dims(num_valid_points, axis=0),
          grid_cell_size=voxel_grid_cell_size,
          voxels_pad_or_clip_size=None,
          segment_func=tf.math.unsorted_segment_sum))
  voxels_density = tf.cast(tf.squeeze(voxels_density, axis=0), dtype=tf.int32)
  segment_ids = tf.squeeze(segment_ids, axis=0)
  points_density = tf.gather(voxels_density, segment_ids)
  valid_points = tf.greater_equal(
      tf.reshape(points_density, [-1]), voxel_density_threshold)
  for key in standard_fields.get_input_point_fields():
    if key != standard_fields.InputDataFields.num_valid_points:
      if key in mesh_inputs:
        mesh_inputs[key] = tf.boolean_mask(mesh_inputs[key], valid_points)
  mesh_inputs[
      standard_fields.InputDataFields.num_valid_points] = tf.math.reduce_sum(
          tf.cast(valid_points, dtype=tf.int32))
  for key in sorted(view_indices_2d_inputs):
    view_indices_2d_inputs[key] = tf.transpose(
        tf.boolean_mask(
            tf.transpose(view_indices_2d_inputs[key], [1, 0, 2]), valid_points),
        [1, 0, 2])
//  [python]
def _get_masked_lm_output(self, inputs, generator, is_training=False):
        """Masked language modeling softmax layer."""
        masked_lm_weights = inputs.masked_lm_weights

        if self._config.uniform_generator:
            logits = tf.zeros(self.disc_config.vocab_size)
            logits_tiled = tf.zeros(
                pretrain_utils.get_shape_list(inputs.masked_lm_ids) +
                [self.disc_config.vocab_size])
            logits_tiled += tf.reshape(logits, [1, 1, self.disc_config.vocab_size])
            logits = logits_tiled
        else:
            outputs = generator(
                input_ids=inputs.input_ids,
                attention_mask=inputs.input_mask,
                token_type_ids=inputs.segment_ids,
                training=is_training)
            logits = outputs[0]
            logits = pretrain_utils.gather_positions(
                logits, inputs.masked_lm_positions)

        oh_labels = tf.one_hot(
            inputs.masked_lm_ids, depth=self.disc_config.vocab_size,
            dtype=tf.float32)

        probs = tf.to_float(tf.nn.softmax(logits))
        log_probs = tf.to_float(tf.nn.log_softmax(logits))
        label_log_probs = -tf.reduce_sum(log_probs * oh_labels, axis=-1)

        numerator = tf.reduce_sum(masked_lm_weights * label_log_probs)
        denominator = tf.reduce_sum(masked_lm_weights) + 1e-6
        loss = numerator / denominator
        preds = tf.argmax(log_probs, axis=-1, output_type=tf.int32)

        MLMOutput = collections.namedtuple(
            "MLMOutput", ["logits", "probs", "loss", "per_example_loss", "preds"])
        return MLMOutput(
            logits=logits, probs=probs, per_example_loss=label_log_probs,
            loss=loss, preds=preds)
//  [python]
def _get_masked_lm_output(self, inputs, generator, is_training=False):
        """Masked language modeling softmax layer."""
        masked_lm_weights = inputs.masked_lm_weights

        if self._config.uniform_generator:
            logits = tf.zeros(self.disc_config.vocab_size)
            logits_tiled = tf.zeros(
                pretrain_utils.get_shape_list(inputs.masked_lm_ids) +
                [self.disc_config.vocab_size])
            logits_tiled += tf.reshape(logits, [1, 1, self.disc_config.vocab_size])
            logits = logits_tiled
        else:
            outputs = generator(
                input_ids=inputs.input_ids,
                attention_mask=inputs.input_mask,
                token_type_ids=inputs.segment_ids,
                training=is_training)
            logits = outputs[0]
            logits = pretrain_utils.gather_positions(
                logits, inputs.masked_lm_positions)

        oh_labels = tf.one_hot(
            inputs.masked_lm_ids, depth=self.disc_config.vocab_size,
            dtype=tf.float32)

        probs = tf.cast(tf.nn.softmax(logits), tf.float32)
        log_probs = tf.cast(tf.nn.log_softmax(logits), tf.float32)
        label_log_probs = -tf.reduce_sum(log_probs * oh_labels, axis=-1)

        numerator = tf.reduce_sum(masked_lm_weights * label_log_probs)
        denominator = tf.reduce_sum(masked_lm_weights) + 1e-6
        loss = numerator / denominator
        preds = tf.argmax(log_probs, axis=-1, output_type=tf.int32)

        MLMOutput = collections.namedtuple(
            "MLMOutput", ["logits", "probs", "loss", "per_example_loss", "preds"])
        return MLMOutput(
            logits=logits, probs=probs, per_example_loss=label_log_probs,
            loss=loss, preds=preds)
//  [python]
def voxelize_point_features(inputs, voxels_pad_or_clip_size,
                            voxel_grid_cell_size, point_feature_keys,
                            point_to_voxel_segment_func,
                            num_frame_to_load=1):
  """Voxelizes the point features.

  Args:
    inputs: A dictionary containing input tensors.
    voxels_pad_or_clip_size: Number of target voxels to pad or clip to. If None,
      it will not perform the voxel padding.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
    point_feature_keys: The keys used to form the voxel features.
    point_to_voxel_segment_func: The function used to aggregate the features of
      the points that fall in the same voxel.
    num_frame_to_load: If greater than 1, offset multiframe point cloud point
      features according to point frame index.
  """
  if len(point_feature_keys) == 1:
    point_features = t.cast(inputs[point_feature_keys[0]], dtype=t.float32)
  else:
    point_features = t.concat(
        [t.cast(inputs[key], dtype=t.float32) for key in point_feature_keys],
        axis=1)
  if num_frame_to_load > 1:
    # Offset feature from different frames to non-overlapping dimensions
    offset_feature_list = []
    for i in range(num_frame_to_load):
      offset_feature_list.append(
          tf.where_v2(
              tf.math.equal(
                  inputs[standard_fields.InputDataFields.point_frame_index], i),
              point_features, 0))
    point_features = t.concat(offset_feature_list, axis=1)

  (voxel_features, voxel_xyz_indices, num_valid_voxels, points_to_voxel_mapping,
   voxel_start_locations) = voxel_utils.pointcloud_to_sparse_voxel_grid(
       points=t.expand_dims(
           inputs[standard_fields.InputDataFields.point_positions], axis=0),
       features=t.expand_dims(point_features, axis=0),
       num_valid_points=t.expand_dims(
           inputs[standard_fields.InputDataFields.num_valid_points], axis=0),
       grid_cell_size=voxel_grid_cell_size,
       voxels_pad_or_clip_size=voxels_pad_or_clip_size,
       segment_func=point_to_voxel_segment_func)
  inputs[standard_fields.InputDataFields.voxel_features] = t.squeeze(
      voxel_features, axis=0)
  inputs[standard_fields.InputDataFields.voxel_xyz_indices] = t.squeeze(
      voxel_xyz_indices, axis=0)
  inputs[standard_fields.InputDataFields.num_valid_voxels] = t.squeeze(
      num_valid_voxels, axis=0)
  inputs[standard_fields.InputDataFields.points_to_voxel_mapping] = t.squeeze(
      points_to_voxel_mapping, axis=0)
  inputs[standard_fields.InputDataFields.voxel_start_locations] = t.squeeze(
      voxel_start_locations, axis=0)
  inputs[standard_fields.InputDataFields.voxel_positions] = (t.cast(
      inputs[standard_fields.InputDataFields.voxel_xyz_indices],
      dtype=t.float32) + 0.5) * t.expand_dims(
          t.convert_to_tensor(voxel_grid_cell_size, dtype=t.float32),
          axis=0) + t.expand_dims(
              inputs[standard_fields.InputDataFields.voxel_start_locations],
              axis=0)
//  [python]
def voxelize_point_features(inputs, voxels_pad_or_clip_size,
                            voxel_grid_cell_size, point_feature_keys,
                            point_to_voxel_segment_func,
                            num_frame_to_load=1):
  """Voxelizes the point features.

  Args:
    inputs: A dictionary containing input tensors.
    voxels_pad_or_clip_size: Number of target voxels to pad or clip to. If None,
      it will not perform the voxel padding.
    voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
      size.
    point_feature_keys: The keys used to form the voxel features.
    point_to_voxel_segment_func: The function used to aggregate the features of
      the points that fall in the same voxel.
    num_frame_to_load: If greater than 1, offset multiframe point cloud point
      features according to point frame index.
  """
  if len(point_feature_keys) == 1:
    point_features = tf.cast(inputs[point_feature_keys[0]], dtype=tf.float32)
  else:
    point_features = tf.concat(
        [tf.cast(inputs[key], dtype=tf.float32) for key in point_feature_keys],
        axis=1)
  if num_frame_to_load > 1:
    # Offset feature from different frames to non-overlapping dimensions
    offset_feature_list = []
    for i in range(num_frame_to_load):
      offset_feature_list.append(
          tf.where_v2(
              tf.math.equal(
                  inputs[standard_fields.InputDataFields.point_frame_index], i),
              point_features, 0))
    point_features = tf.concat(offset_feature_list, axis=1)

  (voxel_features, voxel_xyz_indices, num_valid_voxels, points_to_voxel_mapping,
   voxel_start_locations) = voxel_utils.pointcloud_to_sparse_voxel_grid(
       points=tf.expand_dims(
           inputs[standard_fields.InputDataFields.point_positions], axis=0),
       features=tf.expand_dims(point_features, axis=0),
       num_valid_points=tf.expand_dims(
           inputs[standard_fields.InputDataFields.num_valid_points], axis=0),
       grid_cell_size=voxel_grid_cell_size,
       voxels_pad_or_clip_size=voxels_pad_or_clip_size,
       segment_func=point_to_voxel_segment_func)
  inputs[standard_fields.InputDataFields.voxel_features] = tf.squeeze(
      voxel_features, axis=0)
  inputs[standard_fields.InputDataFields.voxel_xyz_indices] = tf.squeeze(
      voxel_xyz_indices, axis=0)
  inputs[standard_fields.InputDataFields.num_valid_voxels] = tf.squeeze(
      num_valid_voxels, axis=0)
  inputs[standard_fields.InputDataFields.points_to_voxel_mapping] = tf.squeeze(
      points_to_voxel_mapping, axis=0)
  inputs[standard_fields.InputDataFields.voxel_start_locations] = tf.squeeze(
      voxel_start_locations, axis=0)
  inputs[standard_fields.InputDataFields.voxel_positions] = (tf.cast(
      inputs[standard_fields.InputDataFields.voxel_xyz_indices],
      dtype=tf.float32) + 0.5) * tf.expand_dims(
          tf.convert_to_tensor(voxel_grid_cell_size, dtype=tf.float32),
          axis=0) + tf.expand_dims(
              inputs[standard_fields.InputDataFields.voxel_start_locations],
              axis=0)
//  [python]
def call(self, hidden_states, p_mask=None, next_layer_dtype=tf.float32):
        """ Args:
            **p_mask**: (`optional`) ``torch.FloatTensor`` of shape `(batch_size, seq_len)`
                invalid position mask such as query and special symbols (PAD, SEP, CLS)
                1.0 means token should be masked.
        """
        x = tf.squeeze(self.dense(hidden_states), axis=-1,
                       name="squeeze_start_logit_pooler")

        if p_mask is not None:
            x = tf.to_float(x) * (1 - p_mask) - 1e30 * p_mask

        return x
//  [python]
def call(self, hidden_states, p_mask=None, next_layer_dtype=tf.float32):
        """ Args:
            **p_mask**: (`optional`) ``torch.FloatTensor`` of shape `(batch_size, seq_len)`
                invalid position mask such as query and special symbols (PAD, SEP, CLS)
                1.0 means token should be masked.
        """
        x = tf.squeeze(self.dense(hidden_states), axis=-1,
                       name="squeeze_start_logit_pooler")

        if p_mask is not None:
            x = tf.cast(x, tf.float32) * (1 - p_mask) - 1e30 * p_mask

        return x
//  [python]
def _get_discriminator_output(self, inputs, discriminator, labels, is_training=False):
        """Discriminator binary classifier."""

        outputs = discriminator(
            input_ids=inputs.input_ids,
            attention_mask=inputs.input_mask,
            token_type_ids=inputs.segment_ids,
            training=is_training,
        )
        logits = outputs[0]
        weights = t.cast(inputs.input_mask, t.float32)
        labelsf = t.cast(labels, t.float32)
        logits = t.cast(logits, t.float32)
        losses = tf.nn.sigmoid_cross_entropy_with_logits(
            logits=logits, labels=labelsf) * weights
        per_example_loss = (tf.reduce_sum(losses, axis=-1) /
                            (1e-6 + tf.reduce_sum(weights, axis=-1)))
        loss = tf.reduce_sum(losses) / (1e-6 + tf.reduce_sum(weights))
        probs = tf.nn.sigmoid(logits)
        preds = t.cast(tf.round((tf.sign(logits) + 1) / 2), t.int32)
        DiscOutput = collections.namedtuple(
            "DiscOutput", ["loss", "per_example_loss", "probs", "preds",
                           "labels"])
        return DiscOutput(
            loss=loss, per_example_loss=per_example_loss, probs=probs,
            preds=preds, labels=labels,
        )
//  [python]
def _get_discriminator_output(self, inputs, discriminator, labels, is_training=False):
        """Discriminator binary classifier."""

        outputs = discriminator(
            input_ids=inputs.input_ids,
            attention_mask=inputs.input_mask,
            token_type_ids=inputs.segment_ids,
            training=is_training,
        )
        logits = outputs[0]
        weights = tf.cast(inputs.input_mask, tf.float32)
        labelsf = tf.cast(labels, tf.float32)
        logits = tf.cast(logits, tf.float32)
        losses = tf.nn.sigmoid_cross_entropy_with_logits(
            logits=logits, labels=labelsf) * weights
        per_example_loss = (tf.reduce_sum(losses, axis=-1) /
                            (1e-6 + tf.reduce_sum(weights, axis=-1)))
        loss = tf.reduce_sum(losses) / (1e-6 + tf.reduce_sum(weights))
        probs = tf.nn.sigmoid(logits)
        preds = tf.cast(tf.round((tf.sign(logits) + 1) / 2), tf.int32)
        DiscOutput = collections.namedtuple(
            "DiscOutput", ["loss", "per_example_loss", "probs", "preds",
                           "labels"])
        return DiscOutput(
            loss=loss, per_example_loss=per_example_loss, probs=probs,
            preds=preds, labels=labels,
        )
//  [python]
def _masked_softmax(
      self, attention_scores, attention_mask=None, training=None):
    """Normalize the attention scores to probabilities."""
    # `attention_scores` = [B, N, T, S]
    if attention_mask is None:
      return self._softmax(attention_scores)

    # The expand dim happens starting from the `num_heads` dimension,
    # (<batch_dims>, num_heads, <query_attention_dims, key_attention_dims>)
    mask_expansion_axes = [-len(self._attention_axes) * 2 - 1]
    for _ in range(len(attention_scores.shape) - len(attention_mask.shape)):
      attention_mask = array_ops.expand_dims(
          attention_mask, axis=mask_expansion_axes)
    if attention_scores.dtype != attention_mask.dtype:
      attention_mask = tf.cast(attention_mask, attention_scores.dtype)
    attention_mask = self._apply_quantizer(
        'masked_softmax_attention_mask', attention_mask, training)

    # Makes attention_scores >= 0 to avoid masked maximum value be 0.
    attention_scores -= math_ops.reduce_min(
        attention_scores, axis=-1, keepdims=True)
    attention_scores = self._apply_quantizer(
        'masked_softmax_sub1', attention_scores, training)
    attention_scores *= attention_mask
    attention_scores = self._apply_quantizer(
        'masked_softmax_mask1', attention_scores, training)

    # Makes attention_scores <= 0, and become max value be 0.
    attention_scores -= math_ops.reduce_max(
        attention_scores, axis=-1, keepdims=True)
    attention_scores = self._apply_quantizer(
        'masked_softmax_sub2', attention_scores, training)

    # Clip the range of values [-6, 0].
    attention_scores = tf.clip_by_value(
        attention_scores, clip_value_min=-6, clip_value_max=0)
    attention_scores = self._apply_quantizer(
        'masked_softmax_clamp', attention_scores, training)
    # We basically hard-code the to-be-masked-out part have -6.
    # Maximum number is 0. It"s reasonable for 8 bit quantization because
    # e^(0) / e^(-6) < 1/256
    attention_scores *= attention_mask
    attention_scores = self._apply_quantizer(
        'masked_softmax_mask2', attention_scores, training)
    adder = attention_mask - 1.0
    adder = self._apply_quantizer('masked_softmax_adder_sub', adder, training)
    adder *= _MASK_CONSTANT_FOR_INT8_QUANTIZATION
    adder = self._apply_quantizer('masked_softmax_adder_mul', adder, training)
    attention_scores += adder
    attention_scores = self._apply_quantizer(
        'masked_softmax_add', attention_scores, training)

    return self._softmax(attention_scores)
//  [python]
def _get_fake_data(self, inputs, mlm_logits):
        """Sample from the generator to create corrupted input."""
        inputs = pretrain_utils.unmask(inputs)
        disallow = tf.one_hot(
            inputs.masked_lm_ids, depth=self.disc_config.vocab_size,
            dtype=tf.float32) if self._config.disallow_correct else None
        sampled_tokens = tf.stop_gradient(pretrain_utils.sample_from_softmax(
            mlm_logits / self._config.temperature, disallow=disallow))
        sampled_tokids = tf.argmax(sampled_tokens, -1, output_type=tf.int32)
        updated_input_ids, masked = pretrain_utils.scatter_update(
            inputs.input_ids, sampled_tokids, inputs.masked_lm_positions)
        labels = masked * (1 - tf.to_int32(
            tf.equal(updated_input_ids, inputs.input_ids)))
        updated_inputs = pretrain_utils.get_updated_inputs(
            inputs, input_ids=updated_input_ids)
        FakedData = collections.namedtuple("FakedData", [
            "inputs", "is_fake_tokens", "sampled_tokens"])
        return FakedData(inputs=updated_inputs, is_fake_tokens=labels,
                         sampled_tokens=sampled_tokens)