jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_ctc_fsa_fast_bw(targets, seq_lens, blank_idx, label_loop=True):
    """
    See :class:`NativeOp.GetCtcFsaFastBwOp`.
    Generates a FSA with CTC topology. The output format is compatible to :func:`fast_baum_welch`.

    :param tf.Tensor targets: shape (batch,time), int32
    :param tf.Tensor seq_lens: shape (batch), int32
    :param int blank_idx: vocab index of the blank symbol
    :param bool label_loop: True -> normal CTC; False -> RNA-like
    :return: edges, weights, start_end_states;
      edges is (4,num_edges), int32, edges of the graph (from,to,emission_idx,sequence_idx).
      weights is (num_edges,), float32. all zero.
      start_end_states is (2,batch), int32, (start,end) state idx in FSA.
    :rtype: (tf.Tensor,tf.Tensor,tf.Tensor)
    """
    assert targets.get_shape().ndims == 2
    targets_shape = tf.shape(targets)
    targets = tf.cast(targets, tf.int32)
    n_batch = targets_shape[0]
    n_time = targets_shape[1]
    with tf.control_dependencies(
        [
            # The check on the seq lens is important
            # because invalid seq lens might not directly lead to an error here
            # but it might just return an invalid FSA.
            # An invalid FSA can however later cause a crash in the FastBaumWelchOp.
            tf_compat.v1.assert_equal(
                tf.reduce_max(seq_lens),
                n_time,
                data=["get_ctc_fsa_fast_bw seq_lens invalid", seq_lens, n_time, targets_shape],
                summarize=100,
            )
        ]
    ):
        n_edges = n_batch * (5 * (n_time - 1) + 10)  # see op documentation
        weights = tf.zeros((n_edges,))
        maker = OpMaker(OpDescription.from_gen_base(native_op.GetCtcFsaFastBwOp))
        op = maker.make_op()
        edges, start_end_states = op(targets, seq_lens, blank_idx, weights, label_loop)
    return edges, weights, start_end_states
//  [python]
def find_threshold(labels, logits, weights, target_accuracy):

  probabilities = tf.math.softmax(logits)
  weights = weights/tf.reduce_sum(weights)

  entropies = -tf.reduce_sum(probabilities*logits, axis=1)+tf.reduce_logsumexp(logits, axis=1)
  corrects = tf.equal(
      tf.argmax(labels, axis=1),
      tf.argmax(probabilities, axis=1)
    )

  indices_sorted = tf.argsort(entropies, axis=0)
  entropies_sorted = tf.gather(entropies, indices_sorted)
  corrects_sorted = tf.gather(corrects, indices_sorted)
  weights_sorted = tf.gather(weights, indices_sorted)

  numerators_sorted = tf.math.cumsum(weights_sorted*tf.cast(corrects_sorted, probabilities.dtype), axis=0)
  denominators_sorted = tf.math.cumsum(weights_sorted, axis=0)
  accuracies_sorted = numerators_sorted/denominators_sorted

  range = tf.math.cumsum(tf.ones_like(accuracies_sorted, dtype=tf.int64), axis=0)-1  # Subtract one so the range starts as zero
  indices_threshold = tf.where(
    accuracies_sorted > tf.constant(target_accuracy, accuracies_sorted.dtype),
    range,
    tf.zeros_like(range)
  )
  index_threshold = tf.reduce_max(indices_threshold)

  entropy_threshold = tf.gather(entropies_sorted, index_threshold)

  return entropy_threshold
//  [python]
def find_threshold(labels, logits, weights, target_accuracy):

  probabilities = tf.math.softmax(logits)
  weights = weights/tf.reduce_sum(weights)

  entropies = -tf.reduce_sum(probabilities*logits, axis=1)+tf.reduce_logsumexp(logits, axis=1)
  corrects = tf.cast(
    tf.equal(
      tf.argmax(labels, axis=1),
      tf.argmax(probabilities, axis=1)
    ),
    probabilities.dtype
  )

  indices_sorted = tf.argsort(entropies, axis=0)
  entropies_sorted = tf.gather(entropies, indices_sorted)
  corrects_sorted = tf.gather(corrects, indices_sorted)
  weights_sorted = tf.gather(weights, indices_sorted)

  numerators_sorted = tf.math.cumsum(weights_sorted*corrects_sorted, axis=0)
  denominators_sorted = tf.math.cumsum(weights_sorted, axis=0)
  accuracies_sorted = numerators_sorted/denominators_sorted

  range = tf.math.cumsum(tf.ones_like(accuracies_sorted, dtype=tf.int64), axis=0)-1  # Subtract one so the range starts as zero
  indices_threshold = tf.where(
    accuracies_sorted > tf.constant(target_accuracy, accuracies_sorted.dtype),
    range,
    tf.zeros_like(range)
  )
  index_threshold = tf.reduce_max(indices_threshold)

  entropy_threshold = tf.gather(entropies_sorted, index_threshold)

  return entropy_threshold
//  [python]
def _apply(self, words):
        if self.max_distance == 0:
            return tf.identity(words)
        num_words = tf.shape(words)[0]
        offset = tf.random.uniform([num_words], maxval=1) * (self.max_distance + 1)
        offset = tf.to_int32(offset)
        new_pos = tf.argsort(tf.range(num_words) + offset)
        return tf.gather(words, new_pos)
//  [python]
def _apply(self, words):
        if self.max_distance == 0:
            return tf.identity(words)
        num_words = tf.shape(words)[0]
        offset = tf.random.uniform([num_words], maxval=1) * (self.max_distance + 1)
        offset = tf.cast(offset, num_words.dtype)
        new_pos = tf.argsort(tf.range(num_words) + offset)
        return tf.gather(words, new_pos)
//  [python]
def next_edit_distance_reduce(last_row, a, a_n, a_ended, b, b_len, optimal_completion=False, a_blank_idx=None):
    """
    Wraps :class:`NativeOp.NextEditDistanceReduceOp`.

    :param tf.Tensor last_row: 2d (batch,b_time + 1), int32. last edit distances
    :param tf.Tensor a: symbols. 2d (batch|1,n_labels), int32. current.
    :param tf.Tensor a_n: scalar or 1d (batch,), int32. current position
    :param tf.Tensor a_ended: 1d (batch,), int32 (casted from bool, because int32 easier to handle)
    :param tf.Tensor b: symbols. 2d (batch,b_time), int32
    :param tf.Tensor b_len: 1d (batch,), int32
    :param tf.Tensor|int|None a_blank_idx: scalar, int32
    :param bool|tf.Tensor optimal_completion:
    :return: 2d (batch,n_labels), int32, next (unnormalized) (optimal completion) edit distance
    :rtype: tf.Tensor
    """
    optimal_completion = t.cast(t.convert_to_tensor(optimal_completion), t.int32)
    a_ended = t.cast(a_ended, t.int32)
    if a_n.shape.ndims == 0:
        a_n = t.tile(t.expand_dims(a_n, 0), t.shape(a_ended))
    if a_blank_idx is None:
        a_blank_idx = -1
    a_blank_idx = t.cast(a_blank_idx, t.int32)
    maker = OpMaker(OpDescription.from_gen_base(native_op.NextEditDistanceReduceOp))
    op = maker.make_op()
    return op(last_row, a, a_n, a_ended, b, b_len, optimal_completion, a_blank_idx)
//  [python]
def next_edit_distance_reduce(last_row, a, a_n, a_ended, b, b_len, optimal_completion=False, a_blank_idx=None):
    """
    Wraps :class:`NativeOp.NextEditDistanceReduceOp`.

    :param tf.Tensor last_row: 2d (batch,b_time + 1), int32. last edit distances
    :param tf.Tensor a: symbols. 2d (batch|1,n_labels), int32. current.
    :param tf.Tensor a_n: scalar or 1d (batch,), int32. current position
    :param tf.Tensor a_ended: 1d (batch,), int32 (casted from bool, because int32 easier to handle)
    :param tf.Tensor b: symbols. 2d (batch,b_time), int32
    :param tf.Tensor b_len: 1d (batch,), int32
    :param tf.Tensor|int|None a_blank_idx: scalar, int32
    :param bool|tf.Tensor optimal_completion:
    :return: 2d (batch,n_labels), int32, next (unnormalized) (optimal completion) edit distance
    :rtype: tf.Tensor
    """
    optimal_completion = tf.cast(tf.convert_to_tensor(optimal_completion), tf.int32)
    a_ended = tf.cast(a_ended, tf.int32)
    if a_n.shape.ndims == 0:
        a_n = tf.tile(tf.expand_dims(a_n, 0), tf.shape(a_ended))
    if a_blank_idx is None:
        a_blank_idx = -1
    a_blank_idx = tf.cast(a_blank_idx, tf.int32)
    maker = OpMaker(OpDescription.from_gen_base(native_op.NextEditDistanceReduceOp))
    op = maker.make_op()
    return op(last_row, a, a_n, a_ended, b, b_len, optimal_completion, a_blank_idx)
//  [python]
def accuracy_with_threshold(labels, logits, weights, threshold):

  probabilities = tf.math.softmax(logits)
  weights = weights/tf.reduce_sum(weights)

  entropies = -tf.reduce_sum(probabilities*logits, axis=1)+tf.reduce_logsumexp(logits, axis=1)
  corrects = t.cast(
    tf.equal(
      tf.argmax(labels, axis=1),
      tf.argmax(probabilities, axis=1)
    ),
    probabilities.dtype
  )

  masks = tf.where(
    entropies <= threshold,
    tf.ones_like(entropies),
    tf.zeros_like(entropies)
  )

  accuracy_mask = tf.math.divide(
    tf.reduce_sum(weights*masks*corrects),
    tf.reduce_sum(weights*masks)
  )

  return accuracy_mask
//  [python]
def accuracy_with_threshold(labels, logits, weights, threshold):

  probabilities = tf.math.softmax(logits)
  weights = weights/tf.reduce_sum(weights)

  entropies = -tf.reduce_sum(probabilities*logits, axis=1)+tf.reduce_logsumexp(logits, axis=1)
  corrects = tf.cast(
    tf.equal(
      tf.argmax(labels, axis=1),
      tf.argmax(probabilities, axis=1)
    ),
    probabilities.dtype
  )

  masks = tf.where(
    entropies <= threshold,
    tf.ones_like(entropies),
    tf.zeros_like(entropies)
  )

  accuracy_mask = tf.math.divide(
    tf.reduce_sum(weights*masks*corrects),
    tf.reduce_sum(weights*masks)
  )

  return accuracy_mask
//  [python]
def _call(self, tokens, sequence_length, keep_shape, probability):
        input_rank = tokens.shape.rank
        input_is_ragged = isinstance(tokens, tf.RaggedTensor)
        batch_shape = None

        if input_is_ragged:
            if keep_shape:
                raise ValueError("keep_shape is not compatible with a ragged input")
            ragged_tokens = tokens
        elif input_rank == 1:
            tokens = tf.expand_dims(tokens, 0)
            ragged_tokens = tf.RaggedTensor.from_tensor(tokens)
        elif input_rank >= 2:
            if sequence_length is None:
                raise ValueError("sequence_length must be passed for ND dense inputs")
            if input_rank > 2:
                input_shape = misc.shape_list(tokens)
                batch_shape = input_shape[:-1]
                tokens = tf.reshape(tokens, [-1, input_shape[-1]])
                sequence_length = tf.reshape(sequence_length, [-1])
            ragged_tokens = tf.RaggedTensor.from_tensor(tokens, lengths=sequence_length)

        noisy_tokens = tf.map_fn(
            lambda tokens: self._maybe_apply_noise(tokens, probability),
            ragged_tokens,
            fn_output_signature=tf.RaggedTensorSpec(
                shape=[None], dtype=ragged_tokens.dtype, ragged_rank=0
            ),
        )

        if input_is_ragged:
            return noisy_tokens

        new_lengths = tf.to_int32(noisy_tokens.row_lengths())
        noisy_tokens = noisy_tokens.to_tensor(
            shape=tf.shape(tokens) if keep_shape else None
        )
        if input_rank == 1:
            new_lengths = new_lengths[0]
            noisy_tokens = noisy_tokens[0]
        elif batch_shape is not None:
            noisy_tokens = tf.reshape(noisy_tokens, batch_shape + [-1])
            new_lengths = tf.reshape(new_lengths, batch_shape)

        return noisy_tokens, new_lengths
//  [python]
def _call(self, tokens, sequence_length, keep_shape, probability):
        input_rank = tokens.shape.rank
        input_is_ragged = isinstance(tokens, tf.RaggedTensor)
        batch_shape = None

        if input_is_ragged:
            if keep_shape:
                raise ValueError("keep_shape is not compatible with a ragged input")
            ragged_tokens = tokens
        elif input_rank == 1:
            tokens = tf.expand_dims(tokens, 0)
            ragged_tokens = tf.RaggedTensor.from_tensor(tokens)
        elif input_rank >= 2:
            if sequence_length is None:
                raise ValueError("sequence_length must be passed for ND dense inputs")
            if input_rank > 2:
                input_shape = misc.shape_list(tokens)
                batch_shape = input_shape[:-1]
                tokens = tf.reshape(tokens, [-1, input_shape[-1]])
                sequence_length = tf.reshape(sequence_length, [-1])
            ragged_tokens = tf.RaggedTensor.from_tensor(tokens, lengths=sequence_length)

        noisy_tokens = tf.map_fn(
            lambda tokens: self._maybe_apply_noise(tokens, probability),
            ragged_tokens,
            fn_output_signature=tf.RaggedTensorSpec(
                shape=[None], dtype=ragged_tokens.dtype, ragged_rank=0
            ),
        )

        if input_is_ragged:
            return noisy_tokens

        new_lengths = tf.cast(noisy_tokens.row_lengths(), tf.int32)
        noisy_tokens = noisy_tokens.to_tensor(
            shape=tf.shape(tokens) if keep_shape else None
        )
        if input_rank == 1:
            new_lengths = new_lengths[0]
            noisy_tokens = noisy_tokens[0]
        elif batch_shape is not None:
            noisy_tokens = tf.reshape(noisy_tokens, batch_shape + [-1])
            new_lengths = tf.reshape(new_lengths, batch_shape)

        return noisy_tokens, new_lengths
//  [python]
def attention_bias(inputs, mode, inf=-1e9, dtype=None):
    if dtype is None:
        dtype = tf.float32

    if dtype != tf.float32:
        inf = dtype.min

    if mode == 'masking':
        mask = inputs
        ret = (1.0 - mask) * inf
        ret = tf.expand_dims(tf.expand_dims(ret, 1), 1)

    elif mode == 'causal':
        length = inputs
        lower_triangle = tf.linalg.band_part(
            tf.fill([length, length], 1.0), -1, 0)
        ret = inf * (1.0 - lower_triangle)
        ret = tf.reshape(ret, [1, 1, length, length])
    else:
        raise ValueError('Unknown mode %s' % mode)

    return tf.to_float(ret)
//  [python]
def attention_bias(inputs, mode, inf=-1e9, dtype=None):
    if dtype is None:
        dtype = tf.float32

    if dtype != tf.float32:
        inf = dtype.min

    if mode == 'masking':
        mask = inputs
        ret = (1.0 - mask) * inf
        ret = tf.expand_dims(tf.expand_dims(ret, 1), 1)

    elif mode == 'causal':
        length = inputs
        lower_triangle = tf.linalg.band_part(
            tf.fill([length, length], 1.0), -1, 0)
        ret = inf * (1.0 - lower_triangle)
        ret = tf.reshape(ret, [1, 1, length, length])
    else:
        raise ValueError('Unknown mode %s' % mode)

    return tf.cast(ret, dtype)
//  [python]
def call(self, tokens, attentions):
        # remove eos token attentions
        eos_mask = t.cast(tokens != self.eos_idx, attentions.dtype)
        eos_mask = tf.expand_dims(eos_mask, 1) * tf.expand_dims(eos_mask, 2)
        attentions = attentions * eos_mask[:, None, None, :, :]
        attentions = attentions[..., :-1, :-1]
        # remove cls token attentions
        attentions = attentions[..., 1:, 1:]
        batch_size, layers, heads, seqlen, _ = shape_list(attentions)
        attentions = tf.reshape(attentions, (batch_size, layers * heads, seqlen, seqlen))

        # features: batch x channels x tokens x tokens (symmetric)
        attentions = average_product_correct(symmetrize(attentions))
        attentions = tf.transpose(attentions, perm=(0, 2, 3, 1))
        return tf.squeeze(self.regression(attentions), 3)
//  [python]
def call(self, tokens, attentions):
        # remove eos token attentions
        eos_mask = tf.cast(tokens != self.eos_idx, attentions.dtype)
        eos_mask = tf.expand_dims(eos_mask, 1) * tf.expand_dims(eos_mask, 2)
        attentions = attentions * eos_mask[:, None, None, :, :]
        attentions = attentions[..., :-1, :-1]
        # remove cls token attentions
        attentions = attentions[..., 1:, 1:]
        batch_size, layers, heads, seqlen, _ = shape_list(attentions)
        attentions = tf.reshape(attentions, (batch_size, layers * heads, seqlen, seqlen))

        # features: batch x channels x tokens x tokens (symmetric)
        attentions = average_product_correct(symmetrize(attentions))
        attentions = tf.transpose(attentions, perm=(0, 2, 3, 1))
        return tf.squeeze(self.regression(attentions), 3)
//  [python]
def add_timing_signal(x, min_timescale=1.0, max_timescale=1.0e4):
    length = tf.shape(x)[1]
    channels = tf.shape(x)[2]
    position = tf.range(length)
    num_timescales = channels // 2

    log_timescale_increment = \
        (math.log(float(max_timescale) / float(min_timescale)) / (tf.cast(num_timescales, tf.float32) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), tf.float32)
        * -log_timescale_increment)

    scaled_time = \
        tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
    signal = tf.pad(signal, [[0, 0], [0, tf.compat.v1.mod(channels, 2)]])
    signal = tf.reshape(signal, [1, length, channels])

    return x + tf.cast(signal, x.dtype)
//  [python]
def add_timing_signal(x, min_timescale=1.0, max_timescale=1.0e4):
    length = tf.shape(x)[1]
    channels = tf.shape(x)[2]
    position = tf.cast(tf.range(length), tf.float32)
    num_timescales = channels // 2

    log_timescale_increment = \
        (math.log(float(max_timescale) / float(min_timescale)) / (tf.cast(num_timescales, tf.float32) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), tf.float32)
        * -log_timescale_increment)

    scaled_time = \
        tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
    signal = tf.pad(signal, [[0, 0], [0, tf.compat.v1.mod(channels, 2)]])
    signal = tf.reshape(signal, [1, length, channels])

    return x + tf.cast(signal, x.dtype)
//  [python]
def get_learning_rate_decay(learning_rate, global_step, params):
    if params['learning_rate_decay'] in ['linear_warmup_rsqrt_decay', 'noam']:
        step = tf.cast(global_step, dtype=tf.float32)
        warmup_steps = tf.cast(params['warmup_steps'], dtype=tf.float32)
        multiplier = params['hidden_size']**-0.5
        decay = multiplier * tf.minimum((step + 1) * (warmup_steps**-1.5),
                                        (step + 1)**-0.5)
        return learning_rate * decay
    elif params['learning_rate_decay'] == 'piecewise_constant':
        return tf.compat.v1.train.piecewise_constant(
            tf.cast(global_step, dtype=tf.int32),
            params['learning_rate_boundaries'], params['learning_rate_values'])
    elif params['learning_rate_decay'] == 'none':
        return learning_rate
    else:
        raise ValueError('Unknown learning_rate_decay')
//  [python]
def semantic_encoding_graph(self, features, params, name=None):
        hidden_size = params['hidden_size']
        initializer = tf.compat.v1.random_normal_initializer(
            0.0, hidden_size**-0.5, dtype=tf.float32)
        scope = None
        if params['shared_source_target_embedding']:
            vocab_size = params['src_vocab_size']
            scope = 'Shared_Semantic_Embedding'
        elif name == 'source':
            vocab_size = params['src_vocab_size']
            scope = 'Source_Semantic_Embedding'
        elif name == 'target':
            vocab_size = params['trg_vocab_size']
            scope = 'Target_Semantic_Embedding'
        else:
            raise ValueError('error: no right name specified.')

        with tf.compat.v1.variable_scope(scope, reuse=tf.compat.v1.AUTO_REUSE):
            embedding_mat = tf.compat.v1.get_variable(
                'Weights', [vocab_size, hidden_size], initializer=initializer)

        eos_padding = tf.zeros_like(features, dtype=tf.int64)[:, :1]
        input_seq = tf.concat([features, eos_padding], 1)
        input_mask = tf.cast(tf.not_equal(input_seq, 0), dtype=tf.float32)
        shift_input_mask = input_mask[:, :-1]
        shift_input_mask = tf.pad(
            tensor=shift_input_mask,
            paddings=[[0, 0], [1, 0]],
            constant_values=1)

        encoder_input = tf.gather(embedding_mat, tf.cast(input_seq, tf.int32))
        encoder_input = encoder_input * (hidden_size**0.5)
        encoder_input = tf.multiply(encoder_input,
                                    tf.expand_dims(shift_input_mask, 2))

        encoder_self_attention_bias = attention_bias(shift_input_mask,
                                                     'masking')

        if params['residual_dropout'] > 0.0:
            encoder_input = tf.nn.dropout(
                encoder_input, rate=params['residual_dropout'])

        # encode
        encoder_output = transformer_semantic_encoder(
            encoder_input, encoder_self_attention_bias, shift_input_mask,
            params)
        return encoder_output
//  [python]
def encoding_graph(self, features, params):
        src_vocab_size = params['src_vocab_size']
        hidden_size = params['hidden_size']

        initializer = tf.compat.v1.random_normal_initializer(
            0.0, hidden_size**-0.5, dtype=tf.float32)

        if params['shared_source_target_embedding']:
            with tf.compat.v1.variable_scope(
                    'Shared_Embedding', reuse=tf.compat.v1.AUTO_REUSE):
                src_embedding = tf.compat.v1.get_variable(
                    'Weights', [src_vocab_size, hidden_size],
                    initializer=initializer)
        else:
            with tf.compat.v1.variable_scope('Source_Embedding'):
                src_embedding = tf.compat.v1.get_variable(
                    'Weights', [src_vocab_size, hidden_size],
                    initializer=initializer)
        src_bias = tf.compat.v1.get_variable('encoder_input_bias',
                                             [hidden_size])

        eos_padding = tf.zeros_like(features, dtype=tf.int64)[:, :1]
        src_seq = tf.concat([features, eos_padding], 1)
        src_mask = tf.cast(tf.not_equal(src_seq, 0), dtype=tf.float32)
        shift_src_mask = src_mask[:, :-1]
        shift_src_mask = tf.pad(
            tensor=shift_src_mask,
            paddings=[[0, 0], [1, 0]],
            constant_values=1)

        encoder_input = tf.gather(src_embedding, tf.compat.v1.to_int32(src_seq))
        encoder_input = encoder_input * (hidden_size**0.5)
        if params['position_info_type'] == 'absolute':
            encoder_input = add_timing_signal(encoder_input)
        encoder_input = tf.multiply(encoder_input,
                                    tf.expand_dims(shift_src_mask, 2))

        encoder_input = tf.nn.bias_add(encoder_input, src_bias)
        encoder_self_attention_bias = attention_bias(shift_src_mask, 'masking')

        if params['residual_dropout'] > 0.0:
            encoder_input = tf.nn.dropout(
                encoder_input, rate=params['residual_dropout'])

        # encode
        encoder_output = transformer_encoder(encoder_input,
                                             encoder_self_attention_bias,
                                             shift_src_mask, params)
        return encoder_output, encoder_self_attention_bias
//  [python]
def encoding_graph(self, features, params):
        src_vocab_size = params['src_vocab_size']
        hidden_size = params['hidden_size']

        initializer = tf.compat.v1.random_normal_initializer(
            0.0, hidden_size**-0.5, dtype=tf.float32)

        if params['shared_source_target_embedding']:
            with tf.compat.v1.variable_scope(
                    'Shared_Embedding', reuse=tf.compat.v1.AUTO_REUSE):
                src_embedding = tf.compat.v1.get_variable(
                    'Weights', [src_vocab_size, hidden_size],
                    initializer=initializer)
        else:
            with tf.compat.v1.variable_scope('Source_Embedding'):
                src_embedding = tf.compat.v1.get_variable(
                    'Weights', [src_vocab_size, hidden_size],
                    initializer=initializer)
        src_bias = tf.compat.v1.get_variable('encoder_input_bias',
                                             [hidden_size])

        eos_padding = tf.zeros_like(features, dtype=tf.int64)[:, :1]
        src_seq = tf.concat([features, eos_padding], 1)
        src_mask = tf.cast(tf.not_equal(src_seq, 0), dtype=tf.float32)
        shift_src_mask = src_mask[:, :-1]
        shift_src_mask = tf.pad(
            tensor=shift_src_mask,
            paddings=[[0, 0], [1, 0]],
            constant_values=1)

        encoder_input = tf.gather(src_embedding, tf.cast(src_seq, tf.int32))
        encoder_input = encoder_input * (hidden_size**0.5)
        if params['position_info_type'] == 'absolute':
            encoder_input = add_timing_signal(encoder_input)
        encoder_input = tf.multiply(encoder_input,
                                    tf.expand_dims(shift_src_mask, 2))

        encoder_input = tf.nn.bias_add(encoder_input, src_bias)
        encoder_self_attention_bias = attention_bias(shift_src_mask, 'masking')

        if params['residual_dropout'] > 0.0:
            encoder_input = tf.nn.dropout(
                encoder_input, rate=params['residual_dropout'])

        # encode
        encoder_output = transformer_encoder(encoder_input,
                                             encoder_self_attention_bias,
                                             shift_src_mask, params)
        return encoder_output, encoder_self_attention_bias
//  [python]
def predict_contacts(self, tokens, attention_mask):
        attns = self(tokens, attention_mask=attention_mask, return_dict=True, output_attentions=True).attentions
        attns = tf.stack(attns, axis=1)  # Matches the original model layout
        # In the original model, attentions for padding tokens are completely zeroed out.
        # This makes no difference most of the time because the other tokens won't attend to them,
        # but it does for the contact prediction task, which takes attentions as input,
        # so we have to mimic that here.
        attention_mask = tf.cast(attention_mask, attns.dtype)
        attns *= attention_mask[:, None, None, None]
        attns *= attention_mask[:, None, None, :, None]
        return self.contact_head(tokens, attns)
//  [python]
def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.

    Args:
        x: tf.Tensor x:

    Returns: tf.Tensor
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = t.cast(input_ids != padding_idx, t.int64)
    incremental_indices = (t.cumsum(mask, axis=1) + past_key_values_length) * mask
    return incremental_indices + padding_idx
//  [python]
def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.

    Args:
        x: tf.Tensor x:

    Returns: tf.Tensor
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = tf.cast(input_ids != padding_idx, tf.int64)
    incremental_indices = (tf.cumsum(mask, axis=1) + past_key_values_length) * mask
    return incremental_indices + padding_idx
//  [python]
def compute_lower_percentile_means(x, axis, q=50):
  """Computes means of elements less or equal to q percentile along axes.

  Args:
    x: A tensor for input values. Shape = [..., dim_1, ..., dim_len(axis)].
    axis: An integer or a list of integers for percentile reduction axis.
    q: A scalar for percentile.

  Returns:
    A tensor for means of elements less or equal to the 1 percentiles. Shape =
      [...].
  """
  percentiles = tfp.stats.percentile(x, q=q, axis=axis, keepdims=True)
  weights = t.cast(x <= percentiles, dtype=tf.float32)
  return (tf.math.reduce_sum(x * weights, axis=axis) /
          tf.math.reduce_sum(weights, axis=axis))
//  [python]
def compute_lower_percentile_means(x, axis, q=50):
  """Computes means of elements less or equal to q percentile along axes.

  Args:
    x: A tensor for input values. Shape = [..., dim_1, ..., dim_len(axis)].
    axis: An integer or a list of integers for percentile reduction axis.
    q: A scalar for percentile.

  Returns:
    A tensor for means of elements less or equal to the 1 percentiles. Shape =
      [...].
  """
  percentiles = tfp.stats.percentile(x, q=q, axis=axis, keepdims=True)
  weights = tf.cast(x <= percentiles, dtype=tf.float32)
  return (tf.math.reduce_sum(x * weights, axis=axis) /
          tf.math.reduce_sum(weights, axis=axis))
//  [python]
def decoding_graph(self,
                       encoder_output,
                       encoder_self_attention_bias,
                       labels,
                       params={},
                       embedding_augmentation=None):
        trg_vocab_size = params['trg_vocab_size']
        hidden_size = params['hidden_size']

        initializer = tf.compat.v1.random_normal_initializer(
            0.0, hidden_size**-0.5, dtype=tf.float32)

        if params['shared_source_target_embedding']:
            with tf.compat.v1.variable_scope(
                    'Shared_Embedding', reuse=tf.compat.v1.AUTO_REUSE):
                trg_embedding = tf.compat.v1.get_variable(
                    'Weights', [trg_vocab_size, hidden_size],
                    initializer=initializer)
        else:
            with tf.compat.v1.variable_scope('Target_Embedding'):
                trg_embedding = tf.compat.v1.get_variable(
                    'Weights', [trg_vocab_size, hidden_size],
                    initializer=initializer)

        eos_padding = tf.zeros_like(labels, dtype=tf.int64)[:, :1]
        trg_seq = tf.concat([labels, eos_padding], 1)
        trg_mask = tf.cast(tf.not_equal(trg_seq, 0), dtype=tf.float32)
        shift_trg_mask = trg_mask[:, :-1]
        shift_trg_mask = tf.pad(
            tensor=shift_trg_mask,
            paddings=[[0, 0], [1, 0]],
            constant_values=1)

        decoder_input = tf.gather(trg_embedding, tf.cast(trg_seq, tf.int32))

        decoder_input *= hidden_size**0.5
        decoder_self_attention_bias = attention_bias(
            tf.shape(input=decoder_input)[1], 'causal')
        decoder_input = tf.pad(
            tensor=decoder_input, paddings=[[0, 0], [1, 0], [0, 0]])[:, :-1, :]
        if params['position_info_type'] == 'absolute':
            decoder_input = add_timing_signal(decoder_input)

        decoder_input = tf.nn.dropout(
            decoder_input, rate=1 - (1.0 - params['residual_dropout']))

        # training
        decoder_output, attention_weights = transformer_decoder(
            decoder_input,
            encoder_output,
            decoder_self_attention_bias,
            encoder_self_attention_bias,
            states_key=None,
            states_val=None,
            embedding_augmentation=embedding_augmentation,
            params=params)

        logits = self.prediction(decoder_output, params)

        on_value = params['confidence']
        off_value = (1.0 - params['confidence']) / tf.cast(
            trg_vocab_size - 1, dtype=tf.float32)
        soft_targets = tf.one_hot(
            tf.cast(trg_seq, tf.int32),
            depth=trg_vocab_size,
            on_value=on_value,
            off_value=off_value)
        mask = tf.cast(shift_trg_mask, logits.dtype)
        xentropy = tf.nn.softmax_cross_entropy_with_logits(
            logits=logits, labels=tf.stop_gradient(soft_targets)) * mask
        loss = tf.reduce_sum(input_tensor=xentropy) / tf.reduce_sum(
            input_tensor=mask)

        return loss
//  [python]
def MnistInput(mnist_data_file, batch_size, randomize):
  """Create operations to read the MNIST input file.

  Args:
    mnist_data_file: Path of a file containing the MNIST images to process.
    batch_size: size of the mini batches to generate.
    randomize: If true, randomize the dataset.

  Returns:
    images: A tensor with the formatted image data. shape [batch_size, 28*28]
    labels: A tensor with the labels for each image.  shape [batch_size]
  """
  file_queue = tf.train.string_input_producer([mnist_data_file])
  reader = tf.TFRecordReader()
  _, value = reader.read(file_queue)
  example = tf.parse_single_example(
      value,
      features={"image/encoded": tf.FixedLenFeature(shape=(), dtype=tf.string),
                "image/class/label": tf.FixedLenFeature([1], tf.int64)})

  image = tf.image.decode_png(example["image/encoded"], channels=1)
  image = tf.reshape(image, [IMAGE_SIZE * IMAGE_SIZE])
  image = tf.to_float(image)
  image /= 255
  label = tf.to_int32(example["image/class/label"])
  label = tf.reshape(label, [])

  if randomize:
    images, labels = tf.train.shuffle_batch(
        [image, label], batch_size=batch_size,
        capacity=(batch_size * 100),
        min_after_dequeue=(batch_size * 10))
  else:
    images, labels = tf.train.batch([image, label], batch_size=batch_size)

  return images, labels
//  [python]
def MnistInput(mnist_data_file, batch_size, randomize):
  """Create operations to read the MNIST input file.

  Args:
    mnist_data_file: Path of a file containing the MNIST images to process.
    batch_size: size of the mini batches to generate.
    randomize: If true, randomize the dataset.

  Returns:
    images: A tensor with the formatted image data. shape [batch_size, 28*28]
    labels: A tensor with the labels for each image.  shape [batch_size]
  """
  file_queue = tf.train.string_input_producer([mnist_data_file])
  reader = tf.TFRecordReader()
  _, value = reader.read(file_queue)
  example = tf.parse_single_example(
      value,
      features={"image/encoded": tf.FixedLenFeature(shape=(), dtype=tf.string),
                "image/class/label": tf.FixedLenFeature([1], tf.int64)})

  image = tf.cast(tf.image.decode_png(example["image/encoded"], channels=1),
                  tf.float32)
  image = tf.reshape(image, [IMAGE_SIZE * IMAGE_SIZE])
  image /= 255
  label = tf.cast(example["image/class/label"], dtype=tf.int32)
  label = tf.reshape(label, [])

  if randomize:
    images, labels = tf.train.shuffle_batch(
        [image, label], batch_size=batch_size,
        capacity=(batch_size * 100),
        min_after_dequeue=(batch_size * 10))
  else:
    images, labels = tf.train.batch([image, label], batch_size=batch_size)

  return images, labels
//  [python]
def gelu3(x):
    """
    Another version of the GELU, as it is used in PyTorch
    https://github.com/pytorch/pytorch/blob/b7f4b6a6de30116f1b44a08fab9499dd5eb2de7d/test/cpp/api/functional.cpp#L839-L845

    :param tf.Tensor x:
    :rtype: tf.Tensor
    """
    import numpy

    return 0.5 * x * (1 + tf.math.erf(x / tf.cast(numpy.sqrt(2), x.dtype)))
//  [python]
def read_single_example_and_decode(filename_queue):

    # tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)

    # reader = tf.TFRecordReader(options=tfrecord_options)
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)

    features = tf.parse_single_example(
        serialized=serialized_example,
        features={
            'img_name': tf.FixedLenFeature([], tf.string),
            'img_height': tf.FixedLenFeature([], tf.int64),
            'img_width': tf.FixedLenFeature([], tf.int64),
            'img': tf.FixedLenFeature([], tf.string),
            'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
            'num_objects': tf.FixedLenFeature([], tf.int64)
        }
    )
    img_name = features['img_name']
    img_height = tf.to_int32(features['img_height'])
    img_width = tf.to_int32(features['img_width'])
    img = tf.decode_raw(features['img'], tf.uint8)

    img = tf.reshape(img, shape=[img_height, img_width, 3])

    gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
    gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 9])

    num_objects = tf.to_int32(features['num_objects'])
    return img_name, img, gtboxes_and_label, num_objects
//  [python]
def read_single_example_and_decode(filename_queue):

    # tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)

    # reader = tf.TFRecordReader(options=tfrecord_options)
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)

    features = tf.parse_single_example(
        serialized=serialized_example,
        features={
            'img_name': tf.FixedLenFeature([], tf.string),
            'img_height': tf.FixedLenFeature([], tf.int64),
            'img_width': tf.FixedLenFeature([], tf.int64),
            'img': tf.FixedLenFeature([], tf.string),
            'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
            'num_objects': tf.FixedLenFeature([], tf.int64)
        }
    )
    img_name = features['img_name']
    img_height = tf.cast(features['img_height'], tf.int32)
    img_width = tf.cast(features['img_width'], tf.int32)
    img = tf.decode_raw(features['img'], tf.uint8)

    img = tf.reshape(img, shape=[img_height, img_width, 3])

    gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
    gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 9])

    num_objects = tf.cast(features['num_objects'], tf.int32)
    return img_name, img, gtboxes_and_label, num_objects
//  [python]
def variable_scalar_summaries_dict(x, name=None):
    """
    Collects all interesting information about `x`, such as min/max/mean, etc. (all scalars).
    This is used by :func:`variable_summaries`.

    :param tf.Tensor|tf.Variable x:
    :param str name:
    :return: dicth with key -> scalar info, e.g. with "%s_mean" % name -> tf.reduce_mean(x)
    :rtype: dict[str,tf.Tensor]
    """
    if x.dtype == tf.string:
        return {}
    if not name:
        name = get_base_name(x)
    if x.dtype.is_integer:
        x_float = tf.to_float(x)
    else:
        x_float = x
    mean = tf.reduce_mean(x_float)
    stddev = tf.sqrt(tf.reduce_mean(tf.square(x_float - mean)))
    return {
        "%s_mean" % name: mean,
        "%s_stddev" % name: stddev,
        "%s_rms" % name: tf.sqrt(tf.reduce_mean(tf.square(x_float))),
        "%s_l2" % name: tf.sqrt(tf.nn.l2_loss(x_float) * 0.5),
        "%s_max" % name: tf.reduce_max(x),
        "%s_min" % name: tf.reduce_min(x),
    }
//  [python]
def variable_scalar_summaries_dict(x, name=None):
    """
    Collects all interesting information about `x`, such as min/max/mean, etc. (all scalars).
    This is used by :func:`variable_summaries`.

    :param tf.Tensor|tf.Variable x:
    :param str name:
    :return: dicth with key -> scalar info, e.g. with "%s_mean" % name -> tf.reduce_mean(x)
    :rtype: dict[str,tf.Tensor]
    """
    if x.dtype == tf.string:
        return {}
    if not name:
        name = get_base_name(x)
    if x.dtype.is_integer:
        x_float = tf.cast(x, tf.float32)
    else:
        x_float = x
    mean = tf.reduce_mean(x_float)
    stddev = tf.sqrt(tf.reduce_mean(tf.square(x_float - mean)))
    return {
        "%s_mean" % name: mean,
        "%s_stddev" % name: stddev,
        "%s_rms" % name: tf.sqrt(tf.reduce_mean(tf.square(x_float))),
        "%s_l2" % name: tf.sqrt(tf.nn.l2_loss(x_float) * 0.5),
        "%s_max" % name: tf.reduce_max(x),
        "%s_min" % name: tf.reduce_min(x),
    }
//  [python]
def get_loss(end_points, sdf_weight=10., regularization=True, mask_weight = 4.,
             num_sample_points = 2048, FLAGS=None, batch_size=None):

    pred_sdf = end_points['pred_sdf']
    gt_sdf = end_points['ref_sdf']
    ################
    # Compute loss #
    ################
    end_points['losses'] = {}

    ###########sdf
    print ('gt_sdf', gt_sdf.shape)
    print ('pred_sdf', pred_sdf.shape)
    if batch_size is None:
        batch_size = FLAGS.batch_size
    if FLAGS.binary:
        label_sdf = tf.reshape(tf.to_int32(tf.greater(gt_sdf, tf.constant(0.0))),
                               (batch_size, num_sample_points))
        accuracy = tf.reduce_mean(tf.to_float(tf.equal(label_sdf, tf.argmax(pred_sdf, axis=2, output_type=tf.int32))))
        end_points['losses']['accuracy'] = accuracy
        sdf_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=label_sdf, logits=pred_sdf)
        sdf_loss = tf.reduce_mean(sdf_loss) * 100
        end_points['losses']['sdf_loss'] = sdf_loss
    else:
        ############### accuracy
        zero = tf.constant(0, dtype=tf.float32)
        gt_sign = tf.greater(gt_sdf,zero)
        pred_sign = tf.greater(pred_sdf,zero)

        accuracy = tf.reduce_mean(tf.to_float(tf.equal(gt_sign, pred_sign)))
        end_points['losses']['accuracy'] = accuracy
        weight_mask = tf.to_float(tf.less_equal(gt_sdf, tf.constant(0.01))) * mask_weight \
                      + tf.to_float(tf.greater(gt_sdf, tf.constant(0.01)))
        end_points['weighed_mask'] = weight_mask
        sdf_loss = tf.reduce_mean(tf.abs(gt_sdf * sdf_weight - pred_sdf) * weight_mask)
        end_points['losses']['sdf_loss_realvalue'] = tf.reduce_mean(tf.abs(gt_sdf - pred_sdf / sdf_weight))
        sdf_loss = sdf_loss * 1000
    end_points['losses']['sdf_loss'] = sdf_loss
    loss = sdf_loss
    ############### weight decay
    if regularization:
        vgg_regularization_loss = tf.add_n(slim.losses.get_regularization_losses())
        decoder_regularization_loss = tf.add_n(tf.get_collection('regularizer'))
        end_points['losses']['regularization'] = (vgg_regularization_loss + decoder_regularization_loss)
        loss += (vgg_regularization_loss + decoder_regularization_loss)
    end_points['losses']['overall_loss'] = loss
    return loss, end_points
//  [python]
def get_loss(end_points, sdf_weight=10., regularization=True, mask_weight = 4.,
             num_sample_points = 2048, FLAGS=None, batch_size=None):

    pred_sdf = end_points['pred_sdf']
    gt_sdf = end_points['ref_sdf']
    ################
    # Compute loss #
    ################
    end_points['losses'] = {}

    ###########sdf
    print ('gt_sdf', gt_sdf.shape)
    print ('pred_sdf', pred_sdf.shape)
    if batch_size is None:
        batch_size = FLAGS.batch_size
    if FLAGS.binary:
        label_sdf = tf.reshape(tf.cast(tf.math.greater(gt_sdf, tf.constant(0.0)), dtype=tf.int32),
                               (batch_size, num_sample_points))
        accuracy = tf.reduce_mean(tf.cast(tf.equal(label_sdf, tf.argmax(pred_sdf, axis=2, output_type=tf.int32)), dtype=tf.float32))
        end_points['losses']['accuracy'] = accuracy
        sdf_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=label_sdf, logits=pred_sdf)
        sdf_loss = tf.reduce_mean(sdf_loss) * 100
        end_points['losses']['sdf_loss'] = sdf_loss
    else:
        ############### accuracy
        zero = tf.constant(0, dtype=tf.float32)
        gt_sign = tf.greater(gt_sdf,zero)
        pred_sign = tf.greater(pred_sdf,zero)

        accuracy = tf.reduce_mean(tf.cast(tf.equal(gt_sign, pred_sign), dtype=tf.float32))
        end_points['losses']['accuracy'] = accuracy
        weight_mask = tf.to_float(tf.less_equal(gt_sdf, tf.constant(0.01))) * mask_weight \
                      + tf.to_float(tf.greater(gt_sdf, tf.constant(0.01)))
        end_points['weighed_mask'] = weight_mask
        sdf_loss = tf.reduce_mean(tf.abs(gt_sdf * sdf_weight - pred_sdf) * weight_mask)
        end_points['losses']['sdf_loss_realvalue'] = tf.reduce_mean(tf.abs(gt_sdf - pred_sdf / sdf_weight))
        sdf_loss = sdf_loss * 1000
    end_points['losses']['sdf_loss'] = sdf_loss
    loss = sdf_loss
    ############### weight decay
    if regularization:
        vgg_regularization_loss = tf.add_n(slim.losses.get_regularization_losses())
        decoder_regularization_loss = tf.add_n(tf.get_collection('regularizer'))
        end_points['losses']['regularization'] = (vgg_regularization_loss + decoder_regularization_loss)
        loss += (vgg_regularization_loss + decoder_regularization_loss)
    end_points['losses']['overall_loss'] = loss
    return loss, end_points
//  [python]
def unflatten_with_seq_len_mask(x, seq_lens, batch_major=True):
    """
    Basically inverse of :func:`flatten_with_seq_len_mask` and :func:`flatten_with_seq_len_mask_time_major`
    :param tf.Tensor x: shape (time', ...s...s'...) where time' = sum(seq_len) <= batch*time
    :param tf.Tensor seq_lens: shape (batch,) of int32
    :param bool batch_major: if True, the output will be batch major
    :return: tensor of shape (batch,...s..., time, ...s'...) or shape (time,...s...., batch, ...s'...)
    :rtype: tf.Tensor
    """
    with tf.name_scope("unflatten_with_seq_len_mask_time_major"):
        if batch_major:
            mask = sequence_mask(seq_lens)  # shape (batch, time)
        else:
            mask = sequence_mask_time_major(seq_lens)  # shape (time, batch)
        indices = t.cast(tf.where(mask), tf.int32)
        res = tf.scatter_nd(indices, x, shape=tf.concat([tf.shape(mask), tf.shape(x)[1:]], axis=0))
        return res
//  [python]
def unflatten_with_seq_len_mask(x, seq_lens, batch_major=True):
    """
    Basically inverse of :func:`flatten_with_seq_len_mask` and :func:`flatten_with_seq_len_mask_time_major`
    :param tf.Tensor x: shape (time', ...s...s'...) where time' = sum(seq_len) <= batch*time
    :param tf.Tensor seq_lens: shape (batch,) of int32
    :param bool batch_major: if True, the output will be batch major
    :return: tensor of shape (batch,...s..., time, ...s'...) or shape (time,...s...., batch, ...s'...)
    :rtype: tf.Tensor
    """
    with tf.name_scope("unflatten_with_seq_len_mask_time_major"):
        if batch_major:
            mask = sequence_mask(seq_lens)  # shape (batch, time)
        else:
            mask = sequence_mask_time_major(seq_lens)  # shape (time, batch)
        indices = tf.cast(tf.where(mask), tf.int32)
        res = tf.scatter_nd(indices, x, shape=tf.concat([tf.shape(mask), tf.shape(x)[1:]], axis=0))
        return res
//  [python]
def read_and_prepocess_single_img(filename_queue, shortside_len, is_training):

    img_name, img, gtboxes_and_label, num_objects = read_single_example_and_decode(filename_queue)

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

    if is_training:

        if cfgs.RGB2GRAY:
            # img, gtboxes_and_label = image_preprocess.aspect_ratio_jittering(img, gtboxes_and_label)
            img = image_preprocess.random_rgb2gray(img_tensor=img, gtboxes_and_label=gtboxes_and_label)

        if cfgs.IMG_ROTATE:
            # rotate with 0.5 prob. and if rotate, if will random choose a theta from : tf.range(-90, 90+16, delta=15)
            img, gtboxes_and_label = image_preprocess.random_rotate_img(img_tensor=img,
                                                                        gtboxes_and_label=gtboxes_and_label)

        img, gtboxes_and_label, img_h, img_w = image_preprocess.short_side_resize(img_tensor=img,
                                                                                  gtboxes_and_label=gtboxes_and_label,
                                                                                  target_shortside_len=shortside_len,
                                                                                  length_limitation=cfgs.IMG_MAX_LENGTH)

        if cfgs.HORIZONTAL_FLIP:
            img, gtboxes_and_label = image_preprocess.random_flip_left_right(img_tensor=img,
                                                                             gtboxes_and_label=gtboxes_and_label)
        if cfgs.VERTICAL_FLIP:
            img, gtboxes_and_label = image_preprocess.random_flip_up_down(img_tensor=img,
                                                                          gtboxes_and_label=gtboxes_and_label)

    else:
        img, gtboxes_and_label, img_h, img_w = image_preprocess.short_side_resize(img_tensor=img, gtboxes_and_label=gtboxes_and_label,
                                                                                  target_shortside_len=shortside_len,
                                                                                  length_limitation=cfgs.IMG_MAX_LENGTH)
    # gtboxes_and_label = tf.reshape(tf.py_func(filter_small_gt, inp=[gtboxes_and_label], Tout=[tf.float32]), [-1, 9])
    if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']:
        img = img / 255 - tf.constant([[cfgs.PIXEL_MEAN_]])
    else:
        img = img - tf.constant([[cfgs.PIXEL_MEAN]])  # sub pixel mean at last
    return img_name, img, gtboxes_and_label, num_objects, img_h, img_w
//  [python]
def to_int32_64(x):
    """
    :param tf.Tensor x: dtype uint8, int8, int16, int32, int64
    :rtype: tf.Tensor
    :return: dtype int32 or int64
    """
    if x.dtype in [tf.int32, tf.int64]:
        return x
    assert x.dtype in [tf.uint8, tf.int8, tf.uint16, tf.int16]
    return t.cast(x, tf.int32)