jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sequence_loss_sigmoid(logits, targets, weights):
    with tf.name_scope(name="sequence_loss",
                       values=[logits, targets, weights]):
        num_classes = tf.shape(logits)[2]
        logits_flat = tf.reshape(logits, [-1, num_classes])
        targets = tf.to_float(tf.reshape(targets, [-1, num_classes]))
        crossent = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
            labels=targets, logits=logits_flat), axis=1)

        crossent *= tf.reshape(weights, [-1])
        crossent = tf.reduce_sum(crossent)
        total_size = tf.reduce_sum(weights)
        total_size += 1e-12  # to avoid division by 0 for all-0 weights
        crossent /= total_size
    return crossent
//  [python]
def sequence_loss_sigmoid(logits, targets, weights):
    with tf.name_scope(name="sequence_loss",
                       values=[logits, targets, weights]):
        num_classes = tf.shape(logits)[2]
        logits_flat = tf.reshape(logits, [-1, num_classes])
        targets = tf.cast(tf.reshape(targets, [-1, num_classes]), tf.float32)
        crossent = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
            labels=targets, logits=logits_flat), axis=1)

        crossent *= tf.reshape(weights, [-1])
        crossent = tf.reduce_sum(crossent)
        total_size = tf.reduce_sum(weights)
        total_size += 1e-12  # to avoid division by 0 for all-0 weights
        crossent /= total_size
    return crossent
//  [python]
def build_fcn_net(self, inp, use_dice = False):
        def dtype_getter(getter, name, dtype=self.model_dtype, *args, **kwargs):
            var = getter(name, dtype=dtype, *args, **kwargs)
            return var

        def dtype_getter_v(getter, name, dtype=self.model_vdtype, *args, **kwargs):
            var = getter(name, dtype=dtype, *args, **kwargs)
            return var

        with tf.compat.v1.variable_scope("fcn1", custom_getter=dtype_getter, dtype=self.model_dtype):
          with tf.compat.v1.tpu.bfloat16_scope():
            bn1 = tf.compat.v1.layers.batch_normalization(inputs=inp, name='bn1')
            dnn1 = tf.compat.v1.layers.dense(bn1, 200, activation=None, name='f1')
            if use_dice:
                dnn1 = dice(dnn1, name='dice_1', data_type=self.model_dtype)
            else:
                dnn1 = prelu(dnn1, 'prelu1')

            dnn2 = tf.compat.v1.layers.dense(dnn1, 80, activation=None, name='f2')
            if use_dice:
                dnn2 = dice(dnn2, name='dice_2', data_type=self.model_dtype)
            else:
                dnn2 = prelu(dnn2, 'prelu2')
            dnn3 = tf.compat.v1.layers.dense(dnn2, 2, activation=None, name='f3')
          dnn3 = tf.compat.v1.to_float(dnn3)
          self.y_hat = tf.nn.softmax(dnn3) + 0.00000001

        with tf.compat.v1.variable_scope("fcn2", custom_getter=dtype_getter_v, dtype=self.model_vdtype):
          with tf.name_scope("Metrics"):
            #with tf.compat.v1.variable_scope("Metrics", custom_getter=dtype_getter, dtype=self.model_dtype):
                # Cross-entropy loss and optimizer initialization
                ctr_loss = - tf.reduce_mean(tf.math.log(self.y_hat) * self.target_ph)
                self.loss = ctr_loss
                if self.use_negsampling:
                    self.loss += tf.cast(self.aux_loss, self.model_vdtype)
                tf.compat.v1.summary.scalar('loss', self.loss)
                # self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.MomentumOptimizer(learning_rate=self.lr, momentum=0.9).minimize(self.loss)

                # convert sparse optimizer to dense optimizer
                adam_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr)
                gradients = adam_optimizer.compute_gradients(self.loss)
                gradients = self._sparse_to_dense_grads(gradients)
                self.optimizer = adam_optimizer.apply_gradients(gradients)

                # Accuracy metric
                self.accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.round(self.y_hat), self.target_ph), self.model_vdtype))
                tf.compat.v1.summary.scalar('accuracy', self.accuracy)

          self.merged =  tf.compat.v1.summary.merge_all()
//  [python]
def build_fcn_net(self, inp, use_dice = False):
        def dtype_getter(getter, name, dtype=self.model_dtype, *args, **kwargs):
            var = getter(name, dtype=dtype, *args, **kwargs)
            return var

        def dtype_getter_v(getter, name, dtype=self.model_vdtype, *args, **kwargs):
            var = getter(name, dtype=dtype, *args, **kwargs)
            return var

        with tf.compat.v1.variable_scope("fcn1", custom_getter=dtype_getter, dtype=self.model_dtype):
          with tf.compat.v1.tpu.bfloat16_scope():
            bn1 = tf.compat.v1.layers.batch_normalization(inputs=inp, name='bn1')
            dnn1 = tf.compat.v1.layers.dense(bn1, 200, activation=None, name='f1')
            if use_dice:
                dnn1 = dice(dnn1, name='dice_1', data_type=self.model_dtype)
            else:
                dnn1 = prelu(dnn1, 'prelu1')

            dnn2 = tf.compat.v1.layers.dense(dnn1, 80, activation=None, name='f2')
            if use_dice:
                dnn2 = dice(dnn2, name='dice_2', data_type=self.model_dtype)
            else:
                dnn2 = prelu(dnn2, 'prelu2')
            dnn3 = tf.compat.v1.layers.dense(dnn2, 2, activation=None, name='f3')
          dnn3 = tf.cast(dnn3, tf.float32)
          self.y_hat = tf.nn.softmax(dnn3) + 0.00000001

        with tf.compat.v1.variable_scope("fcn2", custom_getter=dtype_getter_v, dtype=self.model_vdtype):
          with tf.name_scope("Metrics"):
            #with tf.compat.v1.variable_scope("Metrics", custom_getter=dtype_getter, dtype=self.model_dtype):
                # Cross-entropy loss and optimizer initialization
                ctr_loss = - tf.reduce_mean(tf.math.log(self.y_hat) * self.target_ph)
                self.loss = ctr_loss
                if self.use_negsampling:
                    self.loss += tf.cast(self.aux_loss, self.model_vdtype)
                tf.compat.v1.summary.scalar('loss', self.loss)
                # self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=self.lr).minimize(self.loss)
                # self.optimizer = tf.compat.v1.train.MomentumOptimizer(learning_rate=self.lr, momentum=0.9).minimize(self.loss)

                # convert sparse optimizer to dense optimizer
                adam_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr)
                gradients = adam_optimizer.compute_gradients(self.loss)
                gradients = self._sparse_to_dense_grads(gradients)
                self.optimizer = adam_optimizer.apply_gradients(gradients)

                # Accuracy metric
                self.accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.round(self.y_hat), self.target_ph), self.model_vdtype))
                tf.compat.v1.summary.scalar('accuracy', self.accuracy)

          self.merged =  tf.compat.v1.summary.merge_all()
//  [python]
def call(self, inputs):
    endpoints = {}
    residual = inputs
    x = self._in_conv(inputs)

    # Top-down
    for i, block in enumerate(self._en_convs):
      x = block(x)
      endpoints[str(i)] = x
      x = self._maxpool(x)

    # Bridge
    for i, block in enumerate(self._bridge_convs):
      x = block(x)

    # Bottom-up
    for i, block in enumerate(self._de_convs):
      dtype = x.dtype
      x = tf.cast(x, tf.float32)
      x = self._upsample(x)
      x = tf.cast(x, dtype)
      x = self._concat([endpoints[str(3-i)], x])
      x = block(x)

    x = self._out_conv(x)
    residual = tf.cast(residual, dtype=x.dtype)
    output = self._sigmoid(x + residual)

    self._output_specs = output.get_shape()
    return output
//  [python]
def process_label(self,inputs,cls_hm,wh_target,weights_):

        inputs= t.cast(inputs, t.float32)

        cls_hm = t.cast(cls_hm, t.float32)/cfg.DATA.use_int8_enlarge

        return inputs,cls_hm,wh_target,weights_
//  [python]
def process_label(self,inputs,cls_hm,wh_target,weights_):

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

        cls_hm = tf.cast(cls_hm, tf.float32)/cfg.DATA.use_int8_enlarge

        return inputs,cls_hm,wh_target,weights_
//  [python]
def ae_loss(tag0,tag1,index,mask):
        '''

        :param tag0: [B,N,C],top left tag
        :param tag1: [B,N,C], bottom right tag
        :param index: [B,M]
        :parma mask: [B,M]
        :return:
        '''
        with tf.name_scope("pull_loss"):
            num = tf.reduce_sum(tf.to_float(mask))+1e-4
            #num = tf.Print(num,["X",num,tf.shape(tag0),tf.shape(tag1),tf.shape(index),tf.shape(mask)],summarize=100)
            tag0 = wmlt.batch_gather(tag0,index[:,:,0])
            tag1 = wmlt.batch_gather(tag1,index[:,:,1])
            tag_mean = (tag0+tag1)/2
            tag0 = tf.pow(tag0-tag_mean,2)/num
            tag0 = tf.reduce_sum(tf.boolean_mask(tag0,mask))
            tag1 = tf.pow(tag1-tag_mean,2)/num
            tag1 = tf.reduce_sum(tf.boolean_mask(tag1,mask))
            #tag0 = tf.Print(tag0,["tag01",tag0,tag1],summarize=100)
            pull = tag0+tag1

        with tf.name_scope("push_loss"):
            neg_index = tfop.make_neg_pair_index(mask)
            push_mask = tf.greater(neg_index,-1)
            neg_index = tf.nn.relu(neg_index)
            num = tf.reduce_sum(tf.to_float(push_mask))+1e-4
            tag0 = wmlt.batch_gather(tag_mean,neg_index[:,:,0])
            tag1 = wmlt.batch_gather(tag_mean,neg_index[:,:,1])
            #tag0 = tf.Print(tag0,["X2",num,tf.shape(tag0),tf.shape(tag1),tf.shape(neg_index),tf.shape(push_mask)],summarize=100)
            tag0 = tf.boolean_mask(tag0,push_mask[...,0])
            tag1 = tf.boolean_mask(tag1,push_mask[...,1])
            #num = tf.Print(num,["X3",num,tf.shape(tag0),tf.shape(tag1),tf.shape(neg_index),tf.shape(push_mask)],summarize=100)
            push = tf.reduce_sum(tf.nn.relu(1-tf.abs(tag0-tag1)))/num
            #push = tf.Print(push,["push",push],summarize=100)

        return pull+push
//  [python]
def ae_loss(tag0,tag1,index,mask):
        '''

        :param tag0: [B,N,C],top left tag
        :param tag1: [B,N,C], bottom right tag
        :param index: [B,M]
        :parma mask: [B,M]
        :return:
        '''
        with tf.name_scope("pull_loss"):
            num = tf.reduce_sum(tf.cast(mask,tf.float32))+1e-4
            #num = tf.Print(num,["X",num,tf.shape(tag0),tf.shape(tag1),tf.shape(index),tf.shape(mask)],summarize=100)
            tag0 = wmlt.batch_gather(tag0,index[:,:,0])
            tag1 = wmlt.batch_gather(tag1,index[:,:,1])
            tag_mean = (tag0+tag1)/2
            tag0 = tf.pow(tag0-tag_mean,2)/num
            tag0 = tf.reduce_sum(tf.boolean_mask(tag0,mask))
            tag1 = tf.pow(tag1-tag_mean,2)/num
            tag1 = tf.reduce_sum(tf.boolean_mask(tag1,mask))
            #tag0 = tf.Print(tag0,["tag01",tag0,tag1],summarize=100)
            pull = tag0+tag1

        with tf.name_scope("push_loss"):
            neg_index = tfop.make_neg_pair_index(mask)
            push_mask = tf.greater(neg_index,-1)
            neg_index = tf.nn.relu(neg_index)
            num = tf.reduce_sum(tf.cast(push_mask,tf.float32))+1e-4
            tag0 = wmlt.batch_gather(tag_mean,neg_index[:,:,0])
            tag1 = wmlt.batch_gather(tag_mean,neg_index[:,:,1])
            #tag0 = tf.Print(tag0,["X2",num,tf.shape(tag0),tf.shape(tag1),tf.shape(neg_index),tf.shape(push_mask)],summarize=100)
            tag0 = tf.boolean_mask(tag0,push_mask[...,0])
            tag1 = tf.boolean_mask(tag1,push_mask[...,1])
            #num = tf.Print(num,["X3",num,tf.shape(tag0),tf.shape(tag1),tf.shape(neg_index),tf.shape(push_mask)],summarize=100)
            push = tf.reduce_sum(tf.nn.relu(1-tf.abs(tag0-tag1)))/num
            #push = tf.Print(push,["push",push],summarize=100)

        return pull+push
//  [python]
def _crop_padding(mask, image_info):
  """Crops padded masks to match original image shape.

  Args:
    mask: a padded mask tensor.
    image_info: a tensor that holds information about original and preprocessed
      images.
  Returns:
    cropped and padded masks: tf.Tensor
  """
  image_shape = t.cast(image_info[0, :], t.int32)
  mask = tf.image.crop_to_bounding_box(
      tf.expand_dims(mask, axis=-1), 0, 0,
      image_shape[0], image_shape[1])
  return tf.expand_dims(mask[:, :, 0], axis=0)
//  [python]
def _crop_padding(mask, image_info):
  """Crops padded masks to match original image shape.

  Args:
    mask: a padded mask tensor.
    image_info: a tensor that holds information about original and preprocessed
      images.
  Returns:
    cropped and padded masks: tf.Tensor
  """
  image_shape = tf.cast(image_info[0, :], tf.int32)
  mask = tf.image.crop_to_bounding_box(
      tf.expand_dims(mask, axis=-1), 0, 0,
      image_shape[0], image_shape[1])
  return tf.expand_dims(mask[:, :, 0], axis=0)
//  [python]
def scale(boxlist, y_scale, x_scale):
  """scale box coordinates in x and y dimensions.

  Args:
    boxlist: BoxList holding N boxes
    y_scale: (float) scalar tensor
    x_scale: (float) scalar tensor

  Returns:
    boxlist: BoxList holding N boxes
  """
  with tf.name_scope('Scale'):
    y_scale = tf.cast(y_scale, tf.float32)
    x_scale = tf.cast(x_scale, tf.float32)
    y_min, x_min, y_max, x_max = tf.split(
        value=boxlist.get(), num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxlist = box_list.BoxList(
        tf.concat([y_min, x_min, y_max, x_max], 1))
    return _copy_extra_fields(scaled_boxlist, boxlist)
//  [python]
def get_img_points(sample_pc, trans_mat_right, pred_xyshift, FLAGS, img_size = (137,137)):
    # sample_pc B*N*3
    size_lst = sample_pc.get_shape().as_list()
    homo_pc = tf.concat((sample_pc, tf.ones((size_lst[0], size_lst[1], 1),dtype=np.float32)),axis= -1)
    print("homo_pc.get_shape()", homo_pc.get_shape())
    pc_xyz = tf.matmul(homo_pc, trans_mat_right)
    print("pc_xyz.get_shape()", pc_xyz.get_shape()) # B * N * 3
    pc_xy = tf.to_float(tf.divide(pc_xyz[:,:,:2], tf.expand_dims(pc_xyz[:,:,2], axis = 2)))
    if FLAGS.shift:
        pc_xy = pc_xy + tf.tile(tf.expand_dims(pred_xyshift / 2 * FLAGS.img_h, axis=1), (1,FLAGS.num_points,1))
    mintensor = tf.constant([0.0,0.0], dtype=tf.float32)
    maxtensor = tf.constant([img_size[0]-1, img_size[1]-1], dtype=tf.float32)
    return tf.minimum(maxtensor, tf.maximum(mintensor, pc_xy)), pc_xy
//  [python]
def get_img_points(sample_pc, trans_mat_right, pred_xyshift, FLAGS, img_size = (137,137)):
    # sample_pc B*N*3
    size_lst = sample_pc.get_shape().as_list()
    homo_pc = tf.concat((sample_pc, tf.ones((size_lst[0], size_lst[1], 1),dtype=np.float32)),axis= -1)
    print("homo_pc.get_shape()", homo_pc.get_shape())
    pc_xyz = tf.matmul(homo_pc, trans_mat_right)
    print("pc_xyz.get_shape()", pc_xyz.get_shape()) # B * N * 3
    pc_xy = tf.cast(tf.divide(pc_xyz[:,:,:2], tf.expand_dims(pc_xyz[:,:,2], axis = 2)), dtype=tf.float32)
    if FLAGS.shift:
        pc_xy = pc_xy + tf.tile(tf.expand_dims(pred_xyshift / 2 * FLAGS.img_h, axis=1), (1,FLAGS.num_points,1))
    mintensor = tf.constant([0.0,0.0], dtype=tf.float32)
    maxtensor = tf.constant([img_size[0]-1, img_size[1]-1], dtype=tf.float32)
    return tf.minimum(maxtensor, tf.maximum(mintensor, pc_xy)), pc_xy
//  [python]
def to_absolute_coordinates(boxlist,
                            height,
                            width,
                            check_range=True,
                            maximum_normalized_coordinate=1.1):
  """Converts normalized box coordinates to absolute pixel coordinates.

  This function raises an assertion failed error when the maximum box coordinate
  value is larger than maximum_normalized_coordinate (in which case coordinates
  are already absolute).

  Args:
    boxlist: BoxList with coordinates in range [0, 1].
    height: Maximum value for height of absolute box coordinates.
    width: Maximum value for width of absolute box coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    maximum_normalized_coordinate: Maximum coordinate value to be considered
      as normalized, default to 1.1.

  Returns:
    boxlist with absolute coordinates in terms of the image size.

  """
  with tf.name_scope('ToAbsoluteCoordinates'):
    height = tf.to_float(height)
    width = tf.to_float(width)

    # Ensure range of input boxes is correct.
    if check_range:
      box_maximum = tf.reduce_max(boxlist.get())
      max_assert = tf.Assert(
          tf.greater_equal(maximum_normalized_coordinate, box_maximum),
          ['maximum box coordinate value is larger '
           'than %f: ' % maximum_normalized_coordinate, box_maximum])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(boxlist, height, width)
//  [python]
def to_absolute_coordinates(boxlist,
                            height,
                            width,
                            check_range=True,
                            maximum_normalized_coordinate=1.1):
  """Converts normalized box coordinates to absolute pixel coordinates.

  This function raises an assertion failed error when the maximum box coordinate
  value is larger than maximum_normalized_coordinate (in which case coordinates
  are already absolute).

  Args:
    boxlist: BoxList with coordinates in range [0, 1].
    height: Maximum value for height of absolute box coordinates.
    width: Maximum value for width of absolute box coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    maximum_normalized_coordinate: Maximum coordinate value to be considered
      as normalized, default to 1.1.

  Returns:
    boxlist with absolute coordinates in terms of the image size.

  """
  with tf.name_scope('ToAbsoluteCoordinates'):
    height = tf.cast(height, tf.float32)
    width = tf.cast(width, tf.float32)

    # Ensure range of input boxes is correct.
    if check_range:
      box_maximum = tf.reduce_max(boxlist.get())
      max_assert = tf.Assert(
          tf.greater_equal(maximum_normalized_coordinate, box_maximum),
          ['maximum box coordinate value is larger '
           'than %f: ' % maximum_normalized_coordinate, box_maximum])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(boxlist, height, width)
//  [python]
def _get_fake_data(self, inputs, mlm_logits):
    """Generate corrupted data for discriminator.

    Note it is poosible for sampled token to be the same as the correct one.
    Args:
      inputs: A dict of all inputs, same as the input of `call()` function
      mlm_logits: The generator's output logits

    Returns:
      A dict of generated fake data
    """
    inputs = models.electra_pretrainer.unmask(inputs, duplicate=True)

    # Samples replaced token.
    sampled_tokens = tf.stop_gradient(
        models.electra_pretrainer.sample_from_softmax(
            mlm_logits, disallow=None))
    sampled_tokids = tf.argmax(sampled_tokens, axis=-1, output_type=tf.int32)

    # Prepares input and label for replaced token detection task.
    updated_input_ids, masked = models.electra_pretrainer.scatter_update(
        inputs['input_word_ids'], sampled_tokids, inputs['masked_lm_positions'])
    rtd_labels = masked * (1 - tf.cast(
        tf.equal(updated_input_ids, inputs['input_word_ids']), tf.int32))
    updated_inputs = models.electra_pretrainer.get_updated_inputs(
        inputs, duplicate=True, input_word_ids=updated_input_ids)

    # Samples (candidate_size-1) negatives and concat with true tokens
    disallow = tf.one_hot(
        inputs['masked_lm_ids'], depth=self.vocab_size, dtype=tf.float32)
    sampled_candidates = tf.stop_gradient(
        sample_k_from_softmax(mlm_logits, k=self.candidate_size-1,
                              disallow=disallow))
    true_token_id = tf.expand_dims(inputs['masked_lm_ids'], -1)
    candidate_set = tf.concat([true_token_id, sampled_candidates], -1)

    return {
        'inputs': updated_inputs,
        'is_fake_tokens': rtd_labels,
        'sampled_tokens': sampled_tokens,
        'candidate_set': candidate_set
    }
//  [python]
def _dataset_parser(value, config, is_training):
    """Parse data to a fixed dimension input image and learning targets.
    """
    MAX_NUM_TOKENS = 134 if config['mode'] == 'answer' else 184

    with tf.name_scope('parser'):
        data = _decode_record(value, config)

        # We'll build this one up
        features = {
            'img_id': encode_string(data['img_id'], 24),
            'annot_id': encode_string(data['annot_id'], 24),
        }

        # Normalize img
        image = tf.image.convert_image_dtype(data['image'], dtype=tf.float32)
        image, features['image_info'], _ = \
            resize_and_pad(image, config['image_size'], do_random_scale=False,
                           random_scale_max=1.1, random_scale_min=0.9)

        # bfloat16 and commit the image
        if config.get('use_bfloat16', False):
            tf.logging.info("image -> bfloat")
            image = tf.cast(image, dtype=tf.bfloat16)
        features['images'] = image

        for i in range(4):
            data[f'ctx_choice_{i}'] = tf.concat([tf.constant(VCR_prompt[config['mode']], dtype=tf.int32),
                                                 tf.to_int32(data['ctx']),
                                                 tf.to_int32(data[f'choice_{i}']),
                                                 tf.fill([1], END)], axis=0)

        lm_input = tf.stack([pad_to_fixed_size(data[f'ctx_choice_{i}'],
                                               pad_value=PADDING, output_shape=[MAX_NUM_TOKENS]) for i in range(4)])
        lm_mask = tf.cast(tf.not_equal(lm_input, PADDING), tf.int32)

        features['lm_input'], features['lm_mask'] = lm_input, lm_mask

        ###########
        labels = {
            'lm_targets': data['label'],
        }
        features.update(labels)

        return features, labels
//  [python]
def _dataset_parser(value, config, is_training):
    """Parse data to a fixed dimension input image and learning targets.
    """
    MAX_NUM_TOKENS = 134 if config['mode'] == 'answer' else 184

    with tf.name_scope('parser'):
        data = _decode_record(value, config)

        # We'll build this one up
        features = {
            'img_id': encode_string(data['img_id'], 24),
            'annot_id': encode_string(data['annot_id'], 24),
        }

        # Normalize img
        image = tf.image.convert_image_dtype(data['image'], dtype=tf.float32)
        image, features['image_info'], _ = \
            resize_and_pad(image, config['image_size'], do_random_scale=False,
                           random_scale_max=1.1, random_scale_min=0.9)

        # bfloat16 and commit the image
        if config.get('use_bfloat16', False):
            tf.logging.info("image -> bfloat")
            image = tf.cast(image, dtype=tf.bfloat16)
        features['images'] = image

        for i in range(4):
            data[f'ctx_choice_{i}'] = tf.concat([tf.constant(VCR_prompt[config['mode']], dtype=tf.int32),
                                                 tf.cast(data['ctx'], tf.int32),
                                                 tf.cast(data[f'choice_{i}'], tf.int32),
                                                 tf.fill([1], END)], axis=0)

        lm_input = tf.stack([pad_to_fixed_size(data[f'ctx_choice_{i}'],
                                               pad_value=PADDING, output_shape=[MAX_NUM_TOKENS]) for i in range(4)])
        lm_mask = tf.cast(tf.not_equal(lm_input, PADDING), tf.int32)

        features['lm_input'], features['lm_mask'] = lm_input, lm_mask

        ###########
        labels = {
            'lm_targets': data['label'],
        }
        features.update(labels)

        return features, labels
//  [python]
def postprocess(self, keypoints,wh,max_size):
        """Postprocess outputs of the network.

        Returns:
            boxes: a float tensor with shape [batch_size, N, 4].
            scores: a float tensor with shape [batch_size, N].
            num_boxes: an int tensor with shape [batch_size], it
                represents the number of detections on an image.

            where N = max_boxes.
        """

        def nms(heat, kernel=3):
            hmax = tf.layers.max_pooling2d(heat, kernel, 1, padding='same')
            keep = tf.equal(heat, hmax)
            return heat * tf.cast(keep, tf.float32)

        def topk(hm, K=100):
            batch, height, width, cat = tf.shape(hm)[0], tf.shape(hm)[1], tf.shape(hm)[2], tf.shape(hm)[3]
            # [b,h*w*c]
            scores = tf.reshape(hm, (batch, -1))
            # [b,k]
            topk_scores, topk_inds = tf.nn.top_k(scores, k=K)
            # [b,k]
            topk_clses = topk_inds % cat
            topk_xs = topk_inds // cat % width
            topk_ys = topk_inds // cat // width
            topk_inds = topk_ys * width + topk_xs

            return topk_scores, topk_inds, topk_clses, topk_ys, topk_xs

        def decode(heat, wh, K=100):
            batch, height, width, cat = tf.shape(heat)[0], tf.shape(heat)[1], tf.shape(heat)[2], tf.shape(heat)[3]
            heat = nms(heat)
            scores, inds, clses, ys, xs = topk(heat, K=K)


            xs = tf.expand_dims(tf.cast(xs,tf.float32), axis=-1)
            ys = tf.expand_dims(tf.cast(ys,tf.float32), axis=-1)

            # [b,h*w,2]
            wh = tf.reshape(wh, (batch, -1, tf.shape(wh)[-1]))
            # [b,k,2]
            wh = tf.batch_gather(wh, inds)

            clses = tf.expand_dims(tf.cast(clses, tf.float32), axis=-1)
            scores = tf.expand_dims(scores, axis=-1)

            xmin = xs*cfg.MODEL.global_stride - wh[:,:, 0:1]
            ymin = ys*cfg.MODEL.global_stride - wh[:,:, 1:2]
            xmax = xs*cfg.MODEL.global_stride + wh[:,:, 2:3]
            ymax = ys*cfg.MODEL.global_stride + wh[:,:, 3:4]

            bboxes = tf.concat([xmin, ymin, xmax, ymax], axis=-1)


            # [b,k,6]
            detections = tf.concat([bboxes, scores, clses], axis=-1)
            detections = tf.identity(detections, name='detections')

            # bboxes = tf.identity(bboxes, name='boxes')
            # scores = tf.identity(scores, name='scores')
            # labels = tf.identity(clses, name='labels')  ## no use
            return detections


        decode(keypoints,wh,max_size)
//  [python]
def postprocess(self, keypoints,wh,max_size):
        """Postprocess outputs of the network.

        Returns:
            boxes: a float tensor with shape [batch_size, N, 4].
            scores: a float tensor with shape [batch_size, N].
            num_boxes: an int tensor with shape [batch_size], it
                represents the number of detections on an image.

            where N = max_boxes.
        """

        def nms(heat, kernel=3):
            hmax = tf.layers.max_pooling2d(heat, kernel, 1, padding='same')
            keep = tf.cast(tf.equal(heat, hmax), tf.float32)
            return heat * keep

        def topk(hm, K=100):
            batch, height, width, cat = tf.shape(hm)[0], tf.shape(hm)[1], tf.shape(hm)[2], tf.shape(hm)[3]
            # [b,h*w*c]
            scores = tf.reshape(hm, (batch, -1))
            # [b,k]
            topk_scores, topk_inds = tf.nn.top_k(scores, k=K)
            # [b,k]
            topk_clses = topk_inds % cat
            topk_xs = topk_inds // cat % width
            topk_ys = topk_inds // cat // width
            topk_inds = topk_ys * width + topk_xs

            return topk_scores, topk_inds, topk_clses, topk_ys, topk_xs

        def decode(heat, wh, K=100):
            batch, height, width, cat = tf.shape(heat)[0], tf.shape(heat)[1], tf.shape(heat)[2], tf.shape(heat)[3]
            heat = nms(heat)
            scores, inds, clses, ys, xs = topk(heat, K=K)


            xs = tf.cast(tf.expand_dims(xs, axis=-1),tf.float32)
            ys = tf.cast(tf.expand_dims(ys, axis=-1),tf.float32)

            # [b,h*w,2]
            wh = tf.reshape(wh, (batch, -1, tf.shape(wh)[-1]))
            # [b,k,2]
            wh = tf.batch_gather(wh, inds)

            clses = tf.cast(tf.expand_dims(clses, axis=-1), tf.float32)
            scores = tf.expand_dims(scores, axis=-1)

            xmin = xs*cfg.MODEL.global_stride - wh[:,:, 0:1]
            ymin = ys*cfg.MODEL.global_stride - wh[:,:, 1:2]
            xmax = xs*cfg.MODEL.global_stride + wh[:,:, 2:3]
            ymax = ys*cfg.MODEL.global_stride + wh[:,:, 3:4]

            bboxes = tf.concat([xmin, ymin, xmax, ymax], axis=-1)


            # [b,k,6]
            detections = tf.concat([bboxes, scores, clses], axis=-1)
            detections = tf.identity(detections, name='detections')

            # bboxes = tf.identity(bboxes, name='boxes')
            # scores = tf.identity(scores, name='scores')
            # labels = tf.identity(clses, name='labels')  ## no use
            return detections


        decode(keypoints,wh,max_size)
//  [python]
def build_preprocess(self):
    config = self.config
    image = self.image

    bgr = True  # cv2 load image is bgr
    p_image = tf.expand_dims(image, 0)  # [1,H,W,C]

    with tf.name_scope("image_preprocess"):  # tf.device("/cpu:0"):
      if p_image.dtype.base_dtype != tf.float32:
        p_image = tf.to_float(p_image)

      mean = [0.485, 0.456, 0.406]
      std = [0.229, 0.224, 0.225]

      p_image = p_image*(1.0/255)

      if bgr:
        mean = mean[::-1]
        std = std[::-1]
      image_mean = tf.constant(mean, dtype=tf.float32)
      image_std = tf.constant(std, dtype=tf.float32)
      p_image = (p_image - image_mean) / image_std
      p_image = tf.transpose(p_image, [0, 3, 1, 2])

    self.p_image = p_image
//  [python]
def build_preprocess(self):
    config = self.config
    image = self.image

    bgr = True  # cv2 load image is bgr
    p_image = tf.expand_dims(image, 0)  # [1,H,W,C]

    with tf.name_scope("image_preprocess"):  # tf.device("/cpu:0"):
      if p_image.dtype.base_dtype != tf.float32:
        p_image = tf.cast(p_image, tf.float32)

      mean = [0.485, 0.456, 0.406]
      std = [0.229, 0.224, 0.225]

      p_image = p_image*(1.0/255)

      if bgr:
        mean = mean[::-1]
        std = std[::-1]
      image_mean = tf.constant(mean, dtype=tf.float32)
      image_std = tf.constant(std, dtype=tf.float32)
      p_image = (p_image - image_mean) / image_std
      p_image = tf.transpose(p_image, [0, 3, 1, 2])

    self.p_image = p_image
//  [python]
def build_preprocess(self):
    config = self.config
    image = self.image

    # get feature map anchors first
    # slower if put on cpu # 1.5it/s vs 1.2it/s
    self.multilevel_anchors = []
    with tf.name_scope("fpn_anchors"):#,tf.device("/cpu:0"):
      #fm_h,fm_w = tf.shape(image)[0] // config.anchor_stride,tf.shape(image)[1]
      #// config.anchor_stride
      # all posible anchor box coordinates for a given max_size image,
      # so for 1920 x 1920 image, 1920/16 = 120, so (120,120,NA,4) box, NA is
      #scale*ratio boxes
      self.multilevel_anchors = self.get_all_anchors_fpn()

    bgr = True  # cv2 load image is bgr
    p_image = tf.expand_dims(image, 0)  # [1,H,W,C]

    with tf.name_scope("image_preprocess"):  # tf.device("/cpu:0"):
      if p_image.dtype.base_dtype != tf.float32:
        p_image = tf.to_float(p_image)

      mean = [0.485, 0.456, 0.406]
      std = [0.229, 0.224, 0.225]

      p_image = p_image * (1.0/255)

      if bgr:
        mean = mean[::-1]
        std = std[::-1]
      image_mean = tf.constant(mean, dtype=tf.float32)
      image_std = tf.constant(std, dtype=tf.float32)
      p_image = (p_image - image_mean) / image_std
      p_image = tf.transpose(p_image, [0, 3, 1, 2])

    self.p_image = p_image
//  [python]
def build_preprocess(self):
    config = self.config
    image = self.image

    # get feature map anchors first
    # slower if put on cpu # 1.5it/s vs 1.2it/s
    self.multilevel_anchors = []
    with tf.name_scope("fpn_anchors"):#,tf.device("/cpu:0"):
      #fm_h,fm_w = tf.shape(image)[0] // config.anchor_stride,tf.shape(image)[1]
      #// config.anchor_stride
      # all posible anchor box coordinates for a given max_size image,
      # so for 1920 x 1920 image, 1920/16 = 120, so (120,120,NA,4) box, NA is
      #scale*ratio boxes
      self.multilevel_anchors = self.get_all_anchors_fpn()

    bgr = True  # cv2 load image is bgr
    p_image = tf.expand_dims(image, 0)  # [1,H,W,C]

    with tf.name_scope("image_preprocess"):  # tf.device("/cpu:0"):
      if p_image.dtype.base_dtype != tf.float32:
        p_image = tf.cast(p_image, tf.float32)

      mean = [0.485, 0.456, 0.406]
      std = [0.229, 0.224, 0.225]

      p_image = p_image * (1.0/255)

      if bgr:
        mean = mean[::-1]
        std = std[::-1]
      image_mean = tf.constant(mean, dtype=tf.float32)
      image_std = tf.constant(std, dtype=tf.float32)
      p_image = (p_image - image_mean) / image_std
      p_image = tf.transpose(p_image, [0, 3, 1, 2])

    self.p_image = p_image
//  [python]
def __call__(self, step):
    decay_exp = self.decay_exp
    learning_rate = self.learning_rate
    warmup_steps = self.warmup_steps
    decay_steps = self.decay_steps
    decay_start_steps = self.decay_start_steps

    scal = self.batch_size / 2048

    adj_lr = learning_rate * scal
    if warmup_steps == 0:
      return adj_lr

    warmup_lr = step / warmup_steps * adj_lr
    global_step = t.float32(step)
    decay_steps = t.float32(decay_steps)
    decay_start_step = t.float32(decay_start_steps)
    warmup_lr = t.float32(warmup_lr)

    steps_since_decay_start = global_step - decay_start_step
    already_decayed_steps = t.minimum(steps_since_decay_start, decay_steps)
    decay_lr = adj_lr * (
        (decay_steps - already_decayed_steps) / decay_steps)**decay_exp
    decay_lr = t.maximum(0.0001, decay_lr)

    lr = t.where(
        global_step < warmup_steps, warmup_lr,
        t.where(
            t.logical_and(decay_steps > 0, global_step > decay_start_step),
            decay_lr, adj_lr))

    lr = t.maximum(0.01, lr)
    return lr
//  [python]
def __call__(self, step):
    decay_exp = self.decay_exp
    learning_rate = self.learning_rate
    warmup_steps = self.warmup_steps
    decay_steps = self.decay_steps
    decay_start_steps = self.decay_start_steps

    scal = self.batch_size / 2048

    adj_lr = learning_rate * scal
    if warmup_steps == 0:
      return adj_lr

    warmup_lr = step / warmup_steps * adj_lr
    global_step = tf.cast(step, tf.float32)
    decay_steps = tf.cast(decay_steps, tf.float32)
    decay_start_step = tf.cast(decay_start_steps, tf.float32)
    warmup_lr = tf.cast(warmup_lr, tf.float32)

    steps_since_decay_start = global_step - decay_start_step
    already_decayed_steps = tf.minimum(steps_since_decay_start, decay_steps)
    decay_lr = adj_lr * (
        (decay_steps - already_decayed_steps) / decay_steps)**decay_exp
    decay_lr = tf.maximum(0.0001, decay_lr)

    lr = tf.where(
        global_step < warmup_steps, warmup_lr,
        tf.where(
            tf.logical_and(decay_steps > 0, global_step > decay_start_step),
            decay_lr, adj_lr))

    lr = tf.maximum(0.01, lr)
    return lr
//  [python]
def _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len, num_predict):
  """Turn beg and end indices into actual mask."""
  non_func_mask = tf.logical_and(
      tf.not_equal(inputs, SEP_ID), tf.not_equal(inputs, CLS_ID))
  all_indices = tf.where(non_func_mask, tf.range(tgt_len, dtype=tf.int64),
                         tf.constant(-1, shape=[tgt_len], dtype=tf.int64))
  candidate_matrix = t.cast(
      tf.logical_and(all_indices[None, :] >= beg_indices[:, None],
                     all_indices[None, :] < end_indices[:, None]), t.float32)
  cumsum_matrix = tf.reshape(
      tf.cumsum(tf.reshape(candidate_matrix, [-1])), [-1, tgt_len])
  masked_matrix = t.cast(cumsum_matrix <= num_predict, t.float32)
  target_mask = tf.reduce_sum(candidate_matrix * masked_matrix, axis=0)
  is_masked = t.cast(target_mask, t.bool)

  return is_masked, target_mask
//  [python]
def _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len, num_predict):
  """Turn beg and end indices into actual mask."""
  non_func_mask = tf.logical_and(
      tf.not_equal(inputs, SEP_ID), tf.not_equal(inputs, CLS_ID))
  all_indices = tf.where(non_func_mask, tf.range(tgt_len, dtype=tf.int64),
                         tf.constant(-1, shape=[tgt_len], dtype=tf.int64))
  candidate_matrix = tf.cast(
      tf.logical_and(all_indices[None, :] >= beg_indices[:, None],
                     all_indices[None, :] < end_indices[:, None]), tf.float32)
  cumsum_matrix = tf.reshape(
      tf.cumsum(tf.reshape(candidate_matrix, [-1])), [-1, tgt_len])
  masked_matrix = tf.cast(cumsum_matrix <= num_predict, tf.float32)
  target_mask = tf.reduce_sum(candidate_matrix * masked_matrix, axis=0)
  is_masked = tf.cast(target_mask, tf.bool)

  return is_masked, target_mask
//  [python]
def build_preprocess(self):
    config = self.config
    image = self.image

    # get feature map anchors first
    # slower if put on cpu # 1.5it/s vs 1.2it/s
    self.multilevel_anchors = []
    with tf.name_scope("fpn_anchors"):#,tf.device("/cpu:0"):
      #fm_h,fm_w = tf.shape(image)[0] // config.anchor_stride,tf.shape(image)[1]
      #// config.anchor_stride
      # all posible anchor box coordinates for a given max_size image,
      # so for 1920 x 1920 image, 1920/16 = 120, so (120,120,NA,4) box, NA is
      #scale*ratio boxes
      self.multilevel_anchors = self.get_all_anchors_fpn()

    bgr = True  # cv2 load image is bgr
    #p_image = tf.expand_dims(image, 0)  # [1,H,W,C]
    p_image = image

    with tf.name_scope("image_preprocess"):  # tf.device("/cpu:0"):
      if p_image.dtype.base_dtype != tf.float32:
        p_image = tf.to_float(p_image)

      mean = [0.485, 0.456, 0.406]
      std = [0.229, 0.224, 0.225]

      p_image = p_image * (1.0/255)

      if bgr:
        mean = mean[::-1]
        std = std[::-1]
      image_mean = tf.constant(mean, dtype=tf.float32)
      image_std = tf.constant(std, dtype=tf.float32)
      p_image = (p_image - image_mean) / image_std
      p_image = tf.transpose(p_image, [0, 3, 1, 2])

    self.p_image = p_image
//  [python]
def build_preprocess(self):
    config = self.config
    image = self.image

    # get feature map anchors first
    # slower if put on cpu # 1.5it/s vs 1.2it/s
    self.multilevel_anchors = []
    with tf.name_scope("fpn_anchors"):#,tf.device("/cpu:0"):
      #fm_h,fm_w = tf.shape(image)[0] // config.anchor_stride,tf.shape(image)[1]
      #// config.anchor_stride
      # all posible anchor box coordinates for a given max_size image,
      # so for 1920 x 1920 image, 1920/16 = 120, so (120,120,NA,4) box, NA is
      #scale*ratio boxes
      self.multilevel_anchors = self.get_all_anchors_fpn()

    bgr = True  # cv2 load image is bgr
    #p_image = tf.expand_dims(image, 0)  # [1,H,W,C]
    p_image = image

    with tf.name_scope("image_preprocess"):  # tf.device("/cpu:0"):
      if p_image.dtype.base_dtype != tf.float32:
        p_image = tf.cast(p_image, tf.float32)

      mean = [0.485, 0.456, 0.406]
      std = [0.229, 0.224, 0.225]

      p_image = p_image * (1.0/255)

      if bgr:
        mean = mean[::-1]
        std = std[::-1]
      image_mean = tf.constant(mean, dtype=tf.float32)
      image_std = tf.constant(std, dtype=tf.float32)
      p_image = (p_image - image_mean) / image_std
      p_image = tf.transpose(p_image, [0, 3, 1, 2])

    self.p_image = p_image
//  [python]
def call(self, inputs):
    cls = t.cast(self.cls, inputs.dtype)
    cls = cls + tf.zeros_like(inputs[:, 0:1])  # A hacky way to tile.
    x = tf.concat([cls, inputs], axis=1)
    return x
//  [python]
def call(self, inputs):
    cls = tf.cast(self.cls, inputs.dtype)
    cls = cls + tf.zeros_like(inputs[:, 0:1])  # A hacky way to tile.
    x = tf.concat([cls, inputs], axis=1)
    return x
//  [python]
def call(self, inputs, inputs_positions=None):
    del inputs_positions
    pos_embedding = self.pos_embedding
    # inputs.shape is (batch_size, seq_len, emb_dim).
    if inputs.shape[1] != pos_embedding.shape[1]:
      pos_embedding = self._interpolate(
          pos_embedding,
          from_shape=self.posemb_origin_shape,
          to_shape=self.posemb_target_shape)
    pos_embedding = tf.to_float(pos_embedding)

    return inputs + pos_embedding
//  [python]
def call(self, inputs, inputs_positions=None):
    del inputs_positions
    pos_embedding = self.pos_embedding
    # inputs.shape is (batch_size, seq_len, emb_dim).
    if inputs.shape[1] != pos_embedding.shape[1]:
      pos_embedding = self._interpolate(
          pos_embedding,
          from_shape=self.posemb_origin_shape,
          to_shape=self.posemb_target_shape)
    pos_embedding = tf.cast(pos_embedding, inputs.dtype)

    return inputs + pos_embedding
//  [python]
def dummy_ade20k_dataset(image_width, image_height):
  def dummy_data(_):
    dummy_image = tf.zeros((1, image_width, image_height, 3), dtype=tf.float32)
    dummy_masks = tf.zeros((1, image_width, image_height, 1), dtype=tf.float32)
    dummy_valid_masks = tf.to_float(dummy_masks)
    dummy_image_info = tf.zeros((1, 4, 2), dtype=tf.float32)
    return (dummy_image, {
        'masks': dummy_masks,
        'valid_masks': dummy_valid_masks,
        'image_info': dummy_image_info,
    })
  dataset = tf.data.Dataset.range(1)
  dataset = dataset.repeat()
  dataset = dataset.map(
      dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  return dataset
//  [python]
def dummy_ade20k_dataset(image_width, image_height):
  def dummy_data(_):
    dummy_image = tf.zeros((1, image_width, image_height, 3), dtype=tf.float32)
    dummy_masks = tf.zeros((1, image_width, image_height, 1), dtype=tf.float32)
    dummy_valid_masks = tf.cast(dummy_masks, dtype=tf.bool)
    dummy_image_info = tf.zeros((1, 4, 2), dtype=tf.float32)
    return (dummy_image, {
        'masks': dummy_masks,
        'valid_masks': dummy_valid_masks,
        'image_info': dummy_image_info,
    })
  dataset = tf.data.Dataset.range(1)
  dataset = dataset.repeat()
  dataset = dataset.map(
      dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  return dataset
//  [python]
def file_based_input_fn_builder(input_file, name_to_features, batch_size,
                                is_training):
  """Creates an `input_fn` closure."""

  logging.info("Input tfrecord file %s", input_file)

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

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

    return example

  def input_fn():
    """Returns dataset for training/evaluation."""
    num_threads = 8
    if isinstance(input_file, str):
      d = tf.data.TFRecordDataset(input_file)
      # For training, we want a lot of parallel reading and shuffling.
      # For eval, we want no shuffling and parallel reading doesn't matter.
      if is_training:
        d = d.shuffle(2048)
        d = d.repeat()
    else:
      cycle_length = min(num_threads, len(input_file))
      d = tf.data.Dataset.from_tensor_slices(input_file)
      # file level shuffle
      d = d.shuffle(len(input_file)).repeat()

      d = d.interleave(
          tf.data.TFRecordDataset,
          cycle_length=cycle_length)

      if is_training:
        # sample level shuffle
        d = d.shuffle(buffer_size=2048)
    d = d.map(
        lambda record: _decode_record(record, name_to_features),
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
    d = d.batch(batch_size, drop_remainder=is_training)

    # When `input_file` is a path to a single file or a list
    # containing a single path, disable auto sharding so that
    # same input file is sent to all workers.
    if isinstance(input_file, str) or len(input_file) == 1:
      options = tf.data.Options()
      options.experimental_distribute.auto_shard_policy = (
          tf.data.experimental.AutoShardPolicy.OFF)
      d = d.with_options(options)

    d = d.prefetch(tf.data.experimental.AUTOTUNE)
    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, name_to_features, batch_size,
                                is_training):
  """Creates an `input_fn` closure."""

  logging.info("Input tfrecord file %s", input_file)

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

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

    return example

  def input_fn():
    """Returns dataset for training/evaluation."""
    num_threads = 8
    if isinstance(input_file, str):
      d = tf.data.TFRecordDataset(input_file)
      # For training, we want a lot of parallel reading and shuffling.
      # For eval, we want no shuffling and parallel reading doesn't matter.
      if is_training:
        d = d.shuffle(2048)
        d = d.repeat()
    else:
      cycle_length = min(num_threads, len(input_file))
      d = tf.data.Dataset.from_tensor_slices(input_file)
      # file level shuffle
      d = d.shuffle(len(input_file)).repeat()

      d = d.interleave(
          tf.data.TFRecordDataset,
          cycle_length=cycle_length)

      if is_training:
        # sample level shuffle
        d = d.shuffle(buffer_size=2048)
    d = d.map(
        lambda record: _decode_record(record, name_to_features),
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
    d = d.batch(batch_size, drop_remainder=is_training)

    # When `input_file` is a path to a single file or a list
    # containing a single path, disable auto sharding so that
    # same input file is sent to all workers.
    if isinstance(input_file, str) or len(input_file) == 1:
      options = tf.data.Options()
      options.experimental_distribute.auto_shard_policy = (
          tf.data.experimental.AutoShardPolicy.OFF)
      d = d.with_options(options)

    d = d.prefetch(tf.data.experimental.AUTOTUNE)
    return d

  return input_fn
//  [python]
def _token_span_mask(inputs, tgt_len, num_predict, min_num_tokens,
                     max_num_tokens):
  """Sample token spans as prediction targets."""
  mask_alpha = tgt_len / num_predict
  round_to_int = lambda x: t.cast(t.round(x), t.int64)

  # Sample span lengths from a zipf distribution
  span_len_seq = np.arange(min_num_tokens, max_num_tokens + 1)
  probs = np.array([1.0 / (i + 1) for i in span_len_seq])

  probs /= np.sum(probs)
  logits = t.constant(np.log(probs), dtype=t.float32)
  span_lens = t.random.categorical(
      logits=logits[None],
      num_samples=num_predict,
      dtype=t.int64,
  )[0] + min_num_tokens

  # Sample the ratio [0.0, 1.0) of left context lengths
  span_lens_float = t.cast(span_lens, t.float32)
  left_ratio = t.random.uniform(shape=[num_predict], minval=0.0, maxval=1.0)
  left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)
  left_ctx_len = round_to_int(left_ctx_len)

  # Compute the offset from left start to the right end
  right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

  # Get the actual begin and end indices
  beg_indices = (
      t.cumsum(left_ctx_len) + t.cumsum(right_offset, exclusive=True))
  end_indices = beg_indices + span_lens

  # Remove out of range indices
  valid_idx_mask = end_indices < tgt_len
  beg_indices = t.boolean_mask(beg_indices, valid_idx_mask)
  end_indices = t.boolean_mask(end_indices, valid_idx_mask)

  # Shuffle valid indices
  num_valid = t.cast(t.shape(beg_indices)[0], t.int64)
  order = t.random.shuffle(t.range(num_valid, dtype=t.int64))
  beg_indices = t.gather(beg_indices, order)
  end_indices = t.gather(end_indices, order)

  return _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len,
                           num_predict)