jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: Optional[tf.Tensor] = None,
        token_type_ids: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
        inputs_embeds: Optional[tf.Tensor] = None,
        output_attentions: Optional[tf.Tensor] = None,
        output_hidden_states: Optional[tf.Tensor] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        inputs = input_values_processing(
            func=self.call,
            config=self.config,
            input_values=input_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
            kwargs_call=kwargs,
        )

        hidden_states = self.feature_extractor(
            tf.cast(inputs["input_values"], tf.float32), training=inputs["training"]
        )

        if inputs["attention_mask"] is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(inputs["attention_mask"], -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(hidden_states)[1], dtype=hidden_states.dtype
            )

        hidden_states = self.feature_projection(hidden_states, training=inputs["training"])

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if inputs["training"]:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=inputs["output_attentions"],
            output_hidden_states=inputs["output_hidden_states"],
            return_dict=inputs["return_dict"],
            training=inputs["training"],
        )
        hidden_states = encoder_outputs[0]

        if not inputs["return_dict"]:
            return (hidden_states,) + encoder_outputs[1:]

        return TFBaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )
//  [python]
def testFlat(self, use_tf_function, unroll_fn):
    if use_tf_function:
      unroll_fn = tf.function(unroll_fn)

    initial_state = _, h = self.core.initial_state(self.batch_size)
    input_sequence = tf.random.uniform([self.num_steps, self.batch_size, 1])
    output_sequence, final_state = unroll_fn(self.core, input_sequence,
                                             initial_state)

    self.assertAllClose(
        output_sequence,
        [inputs * (h + t) for t, inputs in enumerate(input_sequence)])
    self.assertAllClose(final_state, (tf.to_float(self.num_steps), h))
//  [python]
def testFlat(self, use_tf_function, unroll_fn):
    if use_tf_function:
      unroll_fn = tf.function(unroll_fn)

    initial_state = _, h = self.core.initial_state(self.batch_size)
    input_sequence = tf.random.uniform([self.num_steps, self.batch_size, 1])
    output_sequence, final_state = unroll_fn(self.core, input_sequence,
                                             initial_state)

    self.assertAllClose(
        output_sequence,
        [inputs * (h + t) for t, inputs in enumerate(input_sequence)])
    self.assertAllClose(final_state, (tf.cast(self.num_steps, tf.float32), h))
//  [python]
def testNestedInputs(self, use_tf_function, unroll_fn):
    if use_tf_function:
      unroll_fn = tf.function(unroll_fn)

    initial_state = _, h = self.core.initial_state(self.batch_size)
    input_sequence = tf.random.uniform([self.num_steps, self.batch_size, 1])
    output_sequence, final_state = unroll_fn(
        lambda inputs, prev_state: self.core(inputs["x"]["y"], prev_state),
        {"x": {
            "y": input_sequence
        }}, initial_state)

    self.assertAllClose(
        output_sequence,
        [inputs * (h + t) for t, inputs in enumerate(input_sequence)])
    self.assertAllClose(final_state, (t.cast(self.num_steps, tf.float32), h))
//  [python]
def testNestedInputs(self, use_tf_function, unroll_fn):
    if use_tf_function:
      unroll_fn = tf.function(unroll_fn)

    initial_state = _, h = self.core.initial_state(self.batch_size)
    input_sequence = tf.random.uniform([self.num_steps, self.batch_size, 1])
    output_sequence, final_state = unroll_fn(
        lambda inputs, prev_state: self.core(inputs["x"]["y"], prev_state),
        {"x": {
            "y": input_sequence
        }}, initial_state)

    self.assertAllClose(
        output_sequence,
        [inputs * (h + t) for t, inputs in enumerate(input_sequence)])
    self.assertAllClose(final_state, (tf.cast(self.num_steps, tf.float32), h))
//  [python]
def testNestedOutputs(self, use_tf_function, unroll_fn):
    if use_tf_function:
      unroll_fn = tf.function(unroll_fn)

    num_replicas = 2
    core = Replicate(self.core, num_replicas)
    initial_state = _, h = core.initial_state(self.batch_size)
    input_sequence = tf.random.uniform([self.num_steps, self.batch_size, 1])
    output_sequence, final_state = unroll_fn(core, input_sequence,
                                             initial_state)

    expected_outputs = [
        inputs * (h + t) for t, inputs in enumerate(input_sequence)
    ]
    self.assertAllClose(output_sequence, (expected_outputs,) * num_replicas)
    self.assertAllClose(final_state, (t.cast(self.num_steps, t.float32), h))
//  [python]
def testNestedOutputs(self, use_tf_function, unroll_fn):
    if use_tf_function:
      unroll_fn = tf.function(unroll_fn)

    num_replicas = 2
    core = Replicate(self.core, num_replicas)
    initial_state = _, h = core.initial_state(self.batch_size)
    input_sequence = tf.random.uniform([self.num_steps, self.batch_size, 1])
    output_sequence, final_state = unroll_fn(core, input_sequence,
                                             initial_state)

    expected_outputs = [
        inputs * (h + t) for t, inputs in enumerate(input_sequence)
    ]
    self.assertAllClose(output_sequence, (expected_outputs,) * num_replicas)
    self.assertAllClose(final_state, (tf.cast(self.num_steps, tf.float32), h))
//  [python]
def build_testing_model(self):
        """
        Similar to the network used for training,
        but without evaluating the corrupted triplet.
        This is used for testing.

        Returns:
            test_predictions: score for edge existence
        """
        log.info('Building the network to be used for testing...')

        # lookup word embeddings
        pred_emb, entity_emb = self._lookup_embeddings()

        # split the input. now there is no corrupted entity in the dataset
        sub, pred, obj = tf.split(t.cast(self.test_triplets, tf.int32), 3, 1)

        # for each term of each sample, select the required embedding
        # and remove the extra dimension caused by the selection
        sub_emb = tf.squeeze(tf.nn.embedding_lookup(entity_emb, sub))
        obj_emb = tf.squeeze(tf.nn.embedding_lookup(entity_emb, obj))
        pred_emb = tf.squeeze(tf.nn.embedding_lookup(pred_emb, pred))

        # calculation of the first layer involves concatenating the three
        # embeddings for each sample and multipling it by the weight vector
        pre_act = tf.add(
            tf.matmul(tf.concat([sub_emb, pred_emb, obj_emb], axis=1), self.weights['C']),
            self.biases['b'])

        # add more layers if necessary
        if self.params['add_layers'] > 0:
            for i in range(1, self.params['add_layers'] + 1):
                post_act = tf.nn.relu(pre_act)

                pre_act = tf.add(
                    tf.matmul(post_act, self.weights['C{}'.format(i)]),
                    self.biases['b{}'.format(i)])

        if self.params['act_function'] == 0:
            # tanh
            log.debug('Using tanh for pre-final layer activation')
            pre_final = tf.tanh(pre_act)
        else:
            # sigmoid
            log.debug('Using sigmoid for pre-final layer activation')
            pre_final = tf.sigmoid(pre_act)

        if self.params['add_layers'] > 0:
            pre_final = tf.nn.dropout(pre_final, self.params['drop_out_percent'])

        out_pre_act = tf.matmul(pre_final, self.weights['B'], name='build_testing_model')

        self.test_predictions = tf.sigmoid(out_pre_act)

        return self.test_predictions
//  [python]
def build_testing_model(self):
        """
        Similar to the network used for training,
        but without evaluating the corrupted triplet.
        This is used for testing.

        Returns:
            test_predictions: score for edge existence
        """
        log.info('Building the network to be used for testing...')

        # lookup word embeddings
        pred_emb, entity_emb = self._lookup_embeddings()

        # split the input. now there is no corrupted entity in the dataset
        sub, pred, obj = tf.split(tf.cast(self.test_triplets, tf.int32), 3, 1)

        # for each term of each sample, select the required embedding
        # and remove the extra dimension caused by the selection
        sub_emb = tf.squeeze(tf.nn.embedding_lookup(entity_emb, sub))
        obj_emb = tf.squeeze(tf.nn.embedding_lookup(entity_emb, obj))
        pred_emb = tf.squeeze(tf.nn.embedding_lookup(pred_emb, pred))

        # calculation of the first layer involves concatenating the three
        # embeddings for each sample and multipling it by the weight vector
        pre_act = tf.add(
            tf.matmul(tf.concat([sub_emb, pred_emb, obj_emb], axis=1), self.weights['C']),
            self.biases['b'])

        # add more layers if necessary
        if self.params['add_layers'] > 0:
            for i in range(1, self.params['add_layers'] + 1):
                post_act = tf.nn.relu(pre_act)

                pre_act = tf.add(
                    tf.matmul(post_act, self.weights['C{}'.format(i)]),
                    self.biases['b{}'.format(i)])

        if self.params['act_function'] == 0:
            # tanh
            log.debug('Using tanh for pre-final layer activation')
            pre_final = tf.tanh(pre_act)
        else:
            # sigmoid
            log.debug('Using sigmoid for pre-final layer activation')
            pre_final = tf.sigmoid(pre_act)

        if self.params['add_layers'] > 0:
            pre_final = tf.nn.dropout(pre_final, self.params['drop_out_percent'])

        out_pre_act = tf.matmul(pre_final, self.weights['B'], name='build_testing_model')

        self.test_predictions = tf.sigmoid(out_pre_act)

        return self.test_predictions
//  [python]
def build_inputs(self, params, input_context=None):
    """Returns tf.data.Dataset for pretraining."""
    # copy from masked_lm.py for testing
    if params.input_path == 'dummy':

      def dummy_data(_):
        dummy_ids = tf.zeros((1, params.seq_length), dtype=tf.int32)
        dummy_lm = tf.zeros((1, params.max_predictions_per_seq), dtype=tf.int32)
        return dict(
            input_word_ids=dummy_ids,
            input_mask=dummy_ids,
            input_type_ids=dummy_ids,
            masked_lm_positions=dummy_lm,
            masked_lm_ids=dummy_lm,
            masked_lm_weights=tf.to_float(dummy_lm),
            next_sentence_labels=tf.zeros((1, 1), dtype=tf.int32))

      dataset = tf.data.Dataset.range(1)
      dataset = dataset.repeat()
      dataset = dataset.map(
          dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return dataset

    return data_loader_factory.get_data_loader(params).load(input_context)
//  [python]
def build_inputs(self, params, input_context=None):
    """Returns tf.data.Dataset for pretraining."""
    # copy from masked_lm.py for testing
    if params.input_path == 'dummy':

      def dummy_data(_):
        dummy_ids = tf.zeros((1, params.seq_length), dtype=tf.int32)
        dummy_lm = tf.zeros((1, params.max_predictions_per_seq), dtype=tf.int32)
        return dict(
            input_word_ids=dummy_ids,
            input_mask=dummy_ids,
            input_type_ids=dummy_ids,
            masked_lm_positions=dummy_lm,
            masked_lm_ids=dummy_lm,
            masked_lm_weights=tf.cast(dummy_lm, dtype=tf.float32),
            next_sentence_labels=tf.zeros((1, 1), dtype=tf.int32))

      dataset = tf.data.Dataset.range(1)
      dataset = dataset.repeat()
      dataset = dataset.map(
          dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return dataset

    return data_loader_factory.get_data_loader(params).load(input_context)
//  [python]
def torch_decay(learning_rate, global_step, decay_rate, name=None):
    '''Adapted from https://github.com/torch/optim/blob/master/adam.lua'''
    if global_step is None:
        raise ValueError("global_step is required for exponential_decay.")
    with tf.name_scope(name, "ExponentialDecay", [learning_rate, global_step, decay_rate]) as name:
        learning_rate = tf.convert_to_tensor(learning_rate, name="learning_rate")
        dtype = learning_rate.dtype
        global_step = tf.to_float(global_step)
        decay_rate = tf.to_float(decay_rate)

        # local clr = lr / (1 + state.t*lrd)
        return learning_rate / (1 + global_step*decay_rate)
//  [python]
def torch_decay(learning_rate, global_step, decay_rate, name=None):
    '''Adapted from https://github.com/torch/optim/blob/master/adam.lua'''
    if global_step is None:
        raise ValueError("global_step is required for exponential_decay.")
    with tf.name_scope(name, "ExponentialDecay", [learning_rate, global_step, decay_rate]) as name:
        learning_rate = tf.convert_to_tensor(learning_rate, name="learning_rate")
        dtype = learning_rate.dtype
        global_step = tf.cast(global_step, dtype)
        decay_rate = tf.cast(decay_rate, dtype)

        # local clr = lr / (1 + state.t*lrd)
        return learning_rate / (1 + global_step*decay_rate)
//  [python]
def add_training_loss(self, final_loss, logits):
    """Computes loss using logits."""
    loss_fn = get_loss_fn(final_loss)  # Get loss function
    task_losses = []
    # label_placeholder of shape (batch_size, n_tasks). Split into n_tasks
    # tensors of shape (batch_size,)
    task_labels = tf.split(
        axis=1, num_or_size_splits=self.n_tasks, value=self.label_placeholder)
    task_weights = tf.split(
        axis=1, num_or_size_splits=self.n_tasks, value=self.weight_placeholder)
    for task in range(self.n_tasks):
      task_label_vector = task_labels[task]
      task_weight_vector = task_weights[task]
      # Convert the labels into one-hot vector encodings.
      one_hot_labels = tf.to_float(
          tf.one_hot(tf.to_int32(tf.squeeze(task_label_vector)), 2))
      # Since we use tf.nn.softmax_cross_entropy_with_logits note that we pass in
      # un-softmaxed logits rather than softmax outputs.
      task_loss = loss_fn(logits[task], one_hot_labels, task_weight_vector)
      task_losses.append(task_loss)
    # It's ok to divide by just the batch_size rather than the number of nonzero
    # examples (effect averages out)
    total_loss = tf.add_n(task_losses)
    total_loss = tf.math.divide(total_loss, self.batch_size)
    return total_loss
//  [python]
def add_training_loss(self, final_loss, logits):
    """Computes loss using logits."""
    loss_fn = get_loss_fn(final_loss)  # Get loss function
    task_losses = []
    # label_placeholder of shape (batch_size, n_tasks). Split into n_tasks
    # tensors of shape (batch_size,)
    task_labels = tf.split(
        axis=1, num_or_size_splits=self.n_tasks, value=self.label_placeholder)
    task_weights = tf.split(
        axis=1, num_or_size_splits=self.n_tasks, value=self.weight_placeholder)
    for task in range(self.n_tasks):
      task_label_vector = task_labels[task]
      task_weight_vector = task_weights[task]
      # Convert the labels into one-hot vector encodings.
      one_hot_labels = tf.cast(
          tf.one_hot(tf.cast(tf.squeeze(task_label_vector), tf.int32), 2),
          tf.float32)
      # Since we use tf.nn.softmax_cross_entropy_with_logits note that we pass in
      # un-softmaxed logits rather than softmax outputs.
      task_loss = loss_fn(logits[task], one_hot_labels, task_weight_vector)
      task_losses.append(task_loss)
    # It's ok to divide by just the batch_size rather than the number of nonzero
    # examples (effect averages out)
    total_loss = tf.add_n(task_losses)
    total_loss = tf.math.divide(total_loss, self.batch_size)
    return total_loss
//  [python]
def test_dtypes_flexible(self):
        """Test that encoder supports variable dtypes and
        yields same result."""

        dt_supported = (tf.float16, tf.float32, tf.float64, tf.int8,
            tf.int32, tf.int64, tf.uint8, tf.uint16, tf.uint32)

        bs = 10
        k = 100
        n = 200

        source = BinarySource()
        enc_ref = LDPC5GEncoder(k, n, dtype=tf.float32)

        u = source([bs, k])
        c_ref = enc_ref(u)

        for dt in dt_supported:
            enc = LDPC5GEncoder(k, n, dtype=dt)
            u_dt = tf.cast(u, dt)
            c = enc(u_dt)

            c_32 = tf.cast(c, tf.float32)

            self.assertTrue(np.array_equal(c_ref.numpy(), c_32.numpy()))
//  [python]
def infer_length(seq, eos=1):
    """
    compute length (including first eos) given output indices and eos code.
    :param seq: tf matrix [time,batch] if time_major else [batch,time]
    :param eos: integer index of end-of-sentence token
    :returns: lengths, int32 vector of [batch_size]
    """
    is_eos = t.cast(t.equal(seq, eos), 'int32')
    count_eos = t.cumsum(is_eos, axis=0, exclusive=True)
    lengths = t.reduce_sum(t.cast(t.equal(count_eos, 0), 'int32'), axis=1)
    return lengths
//  [python]
def infer_length(seq, eos=1):
    """
    compute length (including first eos) given output indices and eos code.
    :param seq: tf matrix [time,batch] if time_major else [batch,time]
    :param eos: integer index of end-of-sentence token
    :returns: lengths, int32 vector of [batch_size]
    """
    is_eos = tf.cast(tf.equal(seq, eos), 'int32')
    count_eos = tf.cumsum(is_eos, axis=0, exclusive=True)
    lengths = tf.reduce_sum(tf.cast(tf.equal(count_eos, 0), 'int32'), axis=1)
    return lengths
//  [python]
def cast_like(x, y):
    """Cast x to y's dtype, if necessary."""
    x = tf.convert_to_tensor(x)
    y = tf.convert_to_tensor(y)

    if x.dtype.base_dtype == y.dtype.base_dtype:
        return x

    cast_x = tf.cast(x, y.dtype)
    if cast_x.device != x.device:
        x_name = "(eager Tensor)"
        try:
            x_name = x.name
        except AttributeError:
            pass
        tf.compat.v1.logging.warning("Cast for %s may induce copy from '%s' to '%s'", x_name, x.device, cast_x.device)
    return cast_x
//  [python]
def test_output_dim(self):
        """ This test combines multiple checks to avoid unnecessary rebuilds of the graph during testing.

        a) Test that output dims are correct (=n)

        b) Test that allzero input leads to  all-zero output.

        c) Test that the systematic part is part of code
        (first 2z bits are punctured!)
        """

        bs = 10
        # (k,n)
        ks = [12, 20, 100, 1234, 2000, 6244, 8448]
        rs = [0.2, 0.34, 0.4, 0.47, 0.7, 0.85, 0.9]
        for k in ks:
            for r in rs:
                n = int(k/r)
                if k>3840 and r<1/3:
                    continue # range is officially not supported
                enc = LDPC5GEncoder(k, n)

                # a) Test for correct dimensions
                u = tf.zeros([bs, k])
                c = enc(u).numpy()
                self.assertTrue(c.shape[-1]==n)

                # b) test for all zero codeword
                c_hat = np.zeros([bs, n])
                self.assertTrue(np.array_equal(c, c_hat))

                # c) Test that systematic part (excluding first 2z pos) is
                # valid
                z = enc._z # access private attribute
                u = tf.cast(tf.random.uniform([bs, k],
                                            0,
                                            2,
                                            tf.int32), tf.float32)
                c = enc(u).numpy()
                self.assertTrue(np.array_equal(u[:,2*z:],c[:,:k-2*z]))
//  [python]
def read_tensor_from_image_file(file_name,
                                input_height=299,
                                input_width=299,
                                input_mean=0,
                                input_std=255):
  input_name = "file_reader"
  output_name = "normalized"
  file_reader = tf.io.read_file(file_name, input_name)
  if file_name.endswith(".png"):
    image_reader = tf.io.decode_png(file_reader, channels=3, name="png_reader")
  elif file_name.endswith(".gif"):
    image_reader = tf.squeeze(tf.io.decode_gif(file_reader, name="gif_reader"))
  elif file_name.endswith(".bmp"):
    image_reader = tf.io.decode_bmp(file_reader, name="bmp_reader")
  else:
    image_reader = tf.io.decode_jpeg(
        file_reader, channels=3, name="jpeg_reader")
  float_caster = tf.image.convert_image_dtype(image_reader, dtype=tf.float32)
  dims_expander = tf.expand_dims(float_caster, 0)
  resized = tf.compat.v1.image.resize_bilinear(
      dims_expander, [input_height, input_width]
  )
  normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
  sess = tf.compat.v1.Session()
  return sess.run(normalized)
//  [python]
def read_tensor_from_image_file(file_name,
                                input_height=299,
                                input_width=299,
                                input_mean=0,
                                input_std=255):
  input_name = "file_reader"
  output_name = "normalized"
  file_reader = tf.io.read_file(file_name, input_name)
  if file_name.endswith(".png"):
    image_reader = tf.io.decode_png(file_reader, channels=3, name="png_reader")
  elif file_name.endswith(".gif"):
    image_reader = tf.squeeze(tf.io.decode_gif(file_reader, name="gif_reader"))
  elif file_name.endswith(".bmp"):
    image_reader = tf.io.decode_bmp(file_reader, name="bmp_reader")
  else:
    image_reader = tf.io.decode_jpeg(
        file_reader, channels=3, name="jpeg_reader")
  float_caster = tf.cast(image_reader, tf.float32)
  dims_expander = tf.expand_dims(float_caster, 0)
  resized = tf.compat.v1.image.resize_bilinear(
      dims_expander, [input_height, input_width]
  )
  normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
  sess = tf.compat.v1.Session()
  return sess.run(normalized)
//  [python]
def maybe_upcast(logits, activation_dtype=None, weight_dtype=None, hparams=None):
    if mixed_precision_is_enabled(activation_dtype, weight_dtype, hparams):
        return t.cast(logits, tf.float32)
    return logits
//  [python]
def maybe_upcast(logits, activation_dtype=None, weight_dtype=None, hparams=None):
    if mixed_precision_is_enabled(activation_dtype, weight_dtype, hparams):
        return tf.cast(logits, tf.float32)
    return logits
//  [python]
def word_dropout(inp, inp_len, voc, dropout=0, method='unk', keep_first=False):
    """
    one function to rule all word dropout methods
    :param inp: tf tensor [batch, ninp]
    :param voc: tfnn vocabulary
    :param dropout: probability of dropping, float scalar in [0, 1]
    :param method: one of several supported methods
        'unk' - replace with voc.unk
        'random_word' - replace with random word except eos or unk
    :return: matrix of same shape as inp with dropout applied to it
    """
    if dropout == 0:
        return inp

    inp_shape = tf.shape(inp)
    border = tf.fill([inp_shape[0], 1], False)
    mask = tf.sequence_mask(inp_len - 1 - keep_first, inp_shape[1] - 1 - keep_first)
    if not keep_first:
        mask = tf.concat((mask, border), axis=1)
    else:
        mask = tf.concat((tf.tile(border, [int(keep_first), 1]), mask, border), axis=1)

    mask = tf.logical_and(mask, tf.random_uniform(inp_shape) < dropout)

    if method == 'unk':
        replacement = tf.fill(inp_shape, t.cast(voc._unk, inp.dtype))
    elif method == 'random_word':
        replacement = tf.random_uniform(inp_shape, minval=max(voc.eos, voc.unk) + 1,
                                        maxval=len(voc), dtype=inp.dtype)
    else:
        raise ValueError("Unknown word dropout method: %r" % method)
    return tf.where(mask, replacement, inp)
//  [python]
def word_dropout(inp, inp_len, voc, dropout=0, method='unk', keep_first=False):
    """
    one function to rule all word dropout methods
    :param inp: tf tensor [batch, ninp]
    :param voc: tfnn vocabulary
    :param dropout: probability of dropping, float scalar in [0, 1]
    :param method: one of several supported methods
        'unk' - replace with voc.unk
        'random_word' - replace with random word except eos or unk
    :return: matrix of same shape as inp with dropout applied to it
    """
    if dropout == 0:
        return inp

    inp_shape = tf.shape(inp)
    border = tf.fill([inp_shape[0], 1], False)
    mask = tf.sequence_mask(inp_len - 1 - keep_first, inp_shape[1] - 1 - keep_first)
    if not keep_first:
        mask = tf.concat((mask, border), axis=1)
    else:
        mask = tf.concat((tf.tile(border, [int(keep_first), 1]), mask, border), axis=1)

    mask = tf.logical_and(mask, tf.random_uniform(inp_shape) < dropout)

    if method == 'unk':
        replacement = tf.fill(inp_shape, tf.cast(voc._unk, inp.dtype))
    elif method == 'random_word':
        replacement = tf.random_uniform(inp_shape, minval=max(voc.eos, voc.unk) + 1,
                                        maxval=len(voc), dtype=inp.dtype)
    else:
        raise ValueError("Unknown word dropout method: %r" % method)
    return tf.where(mask, replacement, inp)
//  [python]
def _batch_normalization(input, is_training=True, scope=None):
    # https://github.com/tensorflow/tensorflow/issues/1122#issuecomment-231917736
    # need to update Op since uodate_collections is not None.
    # tf.cond need
    return tf.cond(tf.cast(is_training, tf.bool),
                   lambda: tf.contrib.layers.batch_norm(input, is_training=True, scope=scope),
                   lambda: tf.contrib.layers.batch_norm(input, is_training=False, reuse=True, scope=scope))
//  [python]
def gather_indices_2d(x, block_shape, block_stride):
    """Getting gather indices."""
    # making an identity matrix kernel
    kernel = tf.eye(block_shape[0] * block_shape[1])
    kernel = reshape_range(kernel, 0, 1, [block_shape[0], block_shape[1], 1])
    # making indices [1, h, w, 1] to appy convs
    x_shape = shape_list(x)
    indices = tf.range(x_shape[2] * x_shape[3])
    indices = tf.reshape(indices, [1, x_shape[2], x_shape[3], 1])
    indices = tf.nn.conv2d(
        t.cast(indices, tf.float32), kernel, strides=[1, block_stride[0], block_stride[1], 1], padding="VALID"
    )
    # making indices [num_blocks, dim] to gather
    dims = shape_list(indices)[:3]
    if all([isinstance(dim, int) for dim in dims]):
        num_blocks = functools.reduce(operator.mul, dims, 1)
    else:
        num_blocks = tf.reduce_prod(dims)
    indices = tf.reshape(indices, [num_blocks, -1])
    return t.cast(indices, tf.int32)
//  [python]
def gather_indices_2d(x, block_shape, block_stride):
    """Getting gather indices."""
    # making an identity matrix kernel
    kernel = tf.eye(block_shape[0] * block_shape[1])
    kernel = reshape_range(kernel, 0, 1, [block_shape[0], block_shape[1], 1])
    # making indices [1, h, w, 1] to appy convs
    x_shape = shape_list(x)
    indices = tf.range(x_shape[2] * x_shape[3])
    indices = tf.reshape(indices, [1, x_shape[2], x_shape[3], 1])
    indices = tf.nn.conv2d(
        tf.cast(indices, tf.float32), kernel, strides=[1, block_stride[0], block_stride[1], 1], padding="VALID"
    )
    # making indices [num_blocks, dim] to gather
    dims = shape_list(indices)[:3]
    if all([isinstance(dim, int) for dim in dims]):
        num_blocks = functools.reduce(operator.mul, dims, 1)
    else:
        num_blocks = tf.reduce_prod(dims)
    indices = tf.reshape(indices, [num_blocks, -1])
    return tf.cast(indices, tf.int32)
//  [python]
def attention(self, tower, transfer):
        '''
        compute the ait attention
        :param tower: the tower, i.e., the q_{t} in the paper, the shape is [None, K]
        :param transfer: the transferred info, i.e., the p_{t-1} in the paper, the shape is [None, K]
        :return: the ait output, the shape is [None, L, K]
        '''
        # (N,L,K)
        inputs = tf.concat([tower[:, None, :], transfer[:, None, :]], axis=1)
        # (N,L,K)*(K,K)->(N,L,K), L=2, K=32 in this.
        Q = tf.tensordot(inputs, self.weights['attention_w1'], axes=1)
        K = tf.tensordot(inputs, self.weights['attention_w2'], axes=1)
        V = tf.tensordot(inputs, self.weights['attention_w3'], axes=1)
        # (N,L)
        a = tf.reduce_sum(tf.multiply(Q, K), axis=-1) / \
            tf.sqrt(tf.cast(inputs.shape[-1], tf.float32))
        a = tf.nn.softmax(a, axis=1)
        # (N,L,K)
        outputs = tf.multiply(a[:, :, None], V)
        return tf.reduce_sum(outputs, axis=1)
//  [python]
def call(self, input_embeddings: tf.Tensor,
           input_mask: tf.Tensor) -> Dict[str, tf.Tensor]:
    batch_size, seq_len, embedding_dim = tf_utils.get_shape_list(
        input_embeddings, expected_rank=3)
    reduced_batch_size = batch_size // self.pack_sequences
    packed_seq_len = self.pack_sequences * seq_len
    packed_embeddings = tf.reshape(
        input_embeddings, [reduced_batch_size, packed_seq_len, embedding_dim])
    input_mask = tf.reshape(input_mask, [reduced_batch_size, packed_seq_len])
    example_ids = 1 + tf.range(self.pack_sequences)
    # Shape: [batch_size, seq_len, pack_sequences].
    example_ids = tf.tile(example_ids[None, :, None],
                          [reduced_batch_size, 1, seq_len])
    example_ids = tf.reshape(example_ids, [reduced_batch_size, packed_seq_len])
    example_ids = tf.where(
        tf.math.equal(input_mask, 0), tf.zeros_like(example_ids), example_ids)
    packing_mask = tf.equal(
        tf.expand_dims(example_ids, 2), tf.expand_dims(example_ids, 1))

    attention_mask = self_attention_mask.get_mask(
        packed_embeddings, input_mask, dtype=tf.bool)

    combined_attention_mask = tf.cast(
        tf.math.logical_and(attention_mask, packing_mask), tf.float32)

    return dict(
        packed_embeddings=packed_embeddings,
        combined_attention_mask=combined_attention_mask)
//  [python]
def call(self, input_embeddings: tf.Tensor,
           input_mask: tf.Tensor) -> Dict[str, tf.Tensor]:
    batch_size, seq_len, embedding_dim = tf_utils.get_shape_list(
        input_embeddings, expected_rank=3)
    reduced_batch_size = batch_size // self.pack_sequences
    packed_seq_len = self.pack_sequences * seq_len
    packed_embeddings = tf.reshape(
        input_embeddings, [reduced_batch_size, packed_seq_len, embedding_dim])
    input_mask = tf.reshape(input_mask, [reduced_batch_size, packed_seq_len])
    example_ids = 1 + tf.range(self.pack_sequences)
    # Shape: [batch_size, seq_len, pack_sequences].
    example_ids = tf.tile(example_ids[None, :, None],
                          [reduced_batch_size, 1, seq_len])
    example_ids = tf.reshape(example_ids, [reduced_batch_size, packed_seq_len])
    example_ids = tf.where(
        tf.math.equal(input_mask, 0), tf.zeros_like(example_ids), example_ids)
    packing_mask = tf.cast(
        tf.equal(
            tf.expand_dims(example_ids, 2), tf.expand_dims(example_ids, 1)),
        dtype=tf.bool)

    attention_mask = self_attention_mask.get_mask(
        packed_embeddings, input_mask, dtype=tf.bool)

    combined_attention_mask = tf.cast(
        tf.math.logical_and(attention_mask, packing_mask), tf.float32)

    return dict(
        packed_embeddings=packed_embeddings,
        combined_attention_mask=combined_attention_mask)
//  [python]
def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : tensor
            The network input.
            For word inputs, please use integer index format, 2D tensor: (batch_size, sentence_length).
        """
        word_embeddings = tf.nn.embedding_lookup(
            params=self.embeddings,
            ids=inputs,
            name='word_embeddings',
        )

        # Zero out embeddings of pad value
        masks = tf.not_equal(inputs, self.pad_value, name='masks')
        word_embeddings *= tf.expand_dims(tf.cast(masks, tf.float32), axis=-1)
        sum_word_embeddings = tf.reduce_sum(input_tensor=word_embeddings, axis=1)

        # Count number of non-padding words in each sentence
        sentence_lengths = tf.math.count_nonzero(
            masks,
            axis=1,
            keepdims=True,
            dtype=tf.float32,
            name='sentence_lengths',
        )

        sentence_embeddings = tf.divide(
            sum_word_embeddings,
            sentence_lengths + 1e-8,  # Add epsilon to avoid dividing by 0
            name='sentence_embeddings'
        )

        outputs = sentence_embeddings

        return outputs
//  [python]
def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : tensor
            The network input.
            For word inputs, please use integer index format, 2D tensor: (batch_size, sentence_length).
        """
        word_embeddings = tf.nn.embedding_lookup(
            params=self.embeddings,
            ids=inputs,
            name='word_embeddings',
        )

        # Zero out embeddings of pad value
        masks = tf.not_equal(inputs, self.pad_value, name='masks')
        word_embeddings *= tf.cast(tf.expand_dims(masks, axis=-1), dtype=tf.float32)
        sum_word_embeddings = tf.reduce_sum(input_tensor=word_embeddings, axis=1)

        # Count number of non-padding words in each sentence
        sentence_lengths = tf.math.count_nonzero(
            masks,
            axis=1,
            keepdims=True,
            dtype=tf.float32,
            name='sentence_lengths',
        )

        sentence_embeddings = tf.divide(
            sum_word_embeddings,
            sentence_lengths + 1e-8,  # Add epsilon to avoid dividing by 0
            name='sentence_embeddings'
        )

        outputs = sentence_embeddings

        return outputs
//  [python]
def call(self, inputs, stride: tf.Tensor):
    if isinstance(inputs, (list, tuple)):
      if len(inputs) == 2:
        input_tensor, attention_mask = inputs
        key_value = None
      elif len(inputs) == 3:
        input_tensor, key_value, attention_mask = inputs
      else:
        raise ValueError(f'Unexpected inputs to {self.__class__} with '
                         f'length at {len(inputs)}.')
    else:
      input_tensor, key_value, attention_mask = (inputs, None, None)

    target_tensor = input_tensor[:, ::stride, :]
    if attention_mask is not None:
      attention_mask = attention_mask[:, ::stride, :]

    if key_value is None:
      key_value = input_tensor

    attention_output = self._attention_layer(
        query=target_tensor, value=key_value, attention_mask=attention_mask)
    attention_output = self._attention_dropout(attention_output)
    attention_output = target_tensor + self._rezero_a * attention_output
    if self._use_layer_norm:
      attention_output = self._attention_layer_norm(attention_output)
    else:
      attention_output = t.float32(attention_output)

    intermediate_output = self._intermediate_dense(attention_output)
    intermediate_output = self._inner_activation_layer(intermediate_output)
    layer_output = self._output_dense(intermediate_output)
    layer_output = self._output_dropout(layer_output)
    layer_output = attention_output + t.float32(self._rezero_a_ffn * layer_output)
    if self._use_layer_norm:
      layer_output = self._output_layer_norm(layer_output)

    return layer_output
//  [python]
def call(self, inputs, stride: tf.Tensor):
    if isinstance(inputs, (list, tuple)):
      if len(inputs) == 2:
        input_tensor, attention_mask = inputs
        key_value = None
      elif len(inputs) == 3:
        input_tensor, key_value, attention_mask = inputs
      else:
        raise ValueError(f'Unexpected inputs to {self.__class__} with '
                         f'length at {len(inputs)}.')
    else:
      input_tensor, key_value, attention_mask = (inputs, None, None)

    target_tensor = input_tensor[:, ::stride, :]
    if attention_mask is not None:
      attention_mask = attention_mask[:, ::stride, :]

    if key_value is None:
      key_value = input_tensor

    attention_output = self._attention_layer(
        query=target_tensor, value=key_value, attention_mask=attention_mask)
    attention_output = self._attention_dropout(attention_output)
    attention_output = target_tensor + self._rezero_a * attention_output
    if self._use_layer_norm:
      attention_output = self._attention_layer_norm(attention_output)
    else:
      attention_output = tf.cast(attention_output, tf.float32)

    intermediate_output = self._intermediate_dense(attention_output)
    intermediate_output = self._inner_activation_layer(intermediate_output)
    layer_output = self._output_dense(intermediate_output)
    layer_output = self._output_dropout(layer_output)
    layer_output = attention_output + tf.cast(self._rezero_a_ffn * layer_output,
                                              tf.float32)
    if self._use_layer_norm:
      layer_output = self._output_layer_norm(layer_output)

    return layer_output
//  [python]
def __call__(self, score_outputs, labels):
    """Computes total RPN detection loss.

    Computes total RPN detection loss including box and score from all levels.

    Args:
      score_outputs: an OrderDict with keys representing levels and values
        representing scores in [batch_size, height, width, num_anchors].
      labels: the dictionary that returned from dataloader that includes
        ground-truth targets.

    Returns:
      rpn_score_loss: a scalar tensor representing total score loss.
    """
    with tf.name_scope('rpn_loss'):
      levels = sorted(score_outputs.keys())

      score_losses = []
      for level in levels:
        score_losses.append(
            self._rpn_score_loss(
                score_outputs[level],
                labels[level],
                normalizer=tf.to_float(
                    tf.shape(score_outputs[level])[0] *
                    self._rpn_batch_size_per_im)))

      # Sums per level losses to total loss.
      return tf.math.add_n(score_losses)
//  [python]
def __call__(self, score_outputs, labels):
    """Computes total RPN detection loss.

    Computes total RPN detection loss including box and score from all levels.

    Args:
      score_outputs: an OrderDict with keys representing levels and values
        representing scores in [batch_size, height, width, num_anchors].
      labels: the dictionary that returned from dataloader that includes
        ground-truth targets.

    Returns:
      rpn_score_loss: a scalar tensor representing total score loss.
    """
    with tf.name_scope('rpn_loss'):
      levels = sorted(score_outputs.keys())

      score_losses = []
      for level in levels:
        score_losses.append(
            self._rpn_score_loss(
                score_outputs[level],
                labels[level],
                normalizer=tf.cast(
                    tf.shape(score_outputs[level])[0] *
                    self._rpn_batch_size_per_im,
                    dtype=score_outputs[level].dtype)))

      # Sums per level losses to total loss.
      return tf.math.add_n(score_losses)
//  [python]
def _rpn_box_loss(self, box_outputs, box_targets, normalizer=1.0):
    """Computes box regression loss."""
    with tf.name_scope('rpn_box_loss'):
      _, height, width, num_anchors_vertices = box_targets.get_shape().as_list()
      # (batch_size, height, width, num_anchors, 4)
      reshaped_box_targets = tf.reshape(
          box_targets, [-1, height, width, num_anchors_vertices // 4, 4])
      # The box is valid if at least one of the ymin, xmin, ymax, ymax is not 0.
      # (batch_size, height, width, num_anchors)
      valid_mask = tf.reduce_any(
          tf.math.abs(reshaped_box_targets) > 1e-6, axis=-1)
      # (batch_size, height, width, num_anchors * 4)
      valid_mask = tf.to_float(
          tf.repeat(valid_mask, 4, axis=-1))
      # (batch_size, height, width, num_anchors * 4, 1)
      box_targets = tf.expand_dims(box_targets, axis=-1)
      # (batch_size, height, width, num_anchors * 4, 1)
      box_outputs = tf.expand_dims(box_outputs, axis=-1)
      box_loss = self._huber_loss(
          box_targets, box_outputs, sample_weight=valid_mask)
      # The loss is normalized by the sum of non-zero weights and additional
      # normalizer provided by the function caller. Using + 0.01 here to avoid
      # division by zero. For each replica, get the sum of non-zero masks. Then
      # get the mean of sums from all replicas. Note there is an extra division
      # by `num_replicas` in train_step(). So it is equivalent to normalizing
      # the box loss by the global sum of non-zero masks.
      replica_context = tf.distribute.get_replica_context()
      valid_mask = tf.reduce_sum(valid_mask)
      valid_mask_mean = replica_context.all_reduce(
          tf.distribute.ReduceOp.MEAN, valid_mask
      )
      box_loss /= normalizer * (valid_mask_mean + 0.01)
      return box_loss
//  [python]
def _rpn_box_loss(self, box_outputs, box_targets, normalizer=1.0):
    """Computes box regression loss."""
    with tf.name_scope('rpn_box_loss'):
      _, height, width, num_anchors_vertices = box_targets.get_shape().as_list()
      # (batch_size, height, width, num_anchors, 4)
      reshaped_box_targets = tf.reshape(
          box_targets, [-1, height, width, num_anchors_vertices // 4, 4])
      # The box is valid if at least one of the ymin, xmin, ymax, ymax is not 0.
      # (batch_size, height, width, num_anchors)
      valid_mask = tf.reduce_any(
          tf.math.abs(reshaped_box_targets) > 1e-6, axis=-1)
      # (batch_size, height, width, num_anchors * 4)
      valid_mask = tf.cast(
          tf.repeat(valid_mask, 4, axis=-1), dtype=box_outputs.dtype)
      # (batch_size, height, width, num_anchors * 4, 1)
      box_targets = tf.expand_dims(box_targets, axis=-1)
      # (batch_size, height, width, num_anchors * 4, 1)
      box_outputs = tf.expand_dims(box_outputs, axis=-1)
      box_loss = self._huber_loss(
          box_targets, box_outputs, sample_weight=valid_mask)
      # The loss is normalized by the sum of non-zero weights and additional
      # normalizer provided by the function caller. Using + 0.01 here to avoid
      # division by zero. For each replica, get the sum of non-zero masks. Then
      # get the mean of sums from all replicas. Note there is an extra division
      # by `num_replicas` in train_step(). So it is equivalent to normalizing
      # the box loss by the global sum of non-zero masks.
      replica_context = tf.distribute.get_replica_context()
      valid_mask = tf.reduce_sum(valid_mask)
      valid_mask_mean = replica_context.all_reduce(
          tf.distribute.ReduceOp.MEAN, valid_mask
      )
      box_loss /= normalizer * (valid_mask_mean + 0.01)
      return box_loss
//  [python]
def aggregate_loss_top_k(self, loss, num_valid_values=None):
    """Aggregate the top-k the greatest loss values.

    Args:
      loss: a float tensor in shape (batch_size, num_boxes) or (batch_size,
        num_boxes, num_classes) which stores the loss values.
      num_valid_values: the number of loss values which are not ignored. The
        default value is None, which means all the loss values are valid.

    Returns:
      A 0-D float which stores the overall loss of the batch.
    """
    loss = tf.reshape(loss, shape=[-1])
    top_k_num = tf.to_int32(
        self._top_k_percent * tf.size(loss, out_type=tf.float32))
    top_k_losses, _ = tf.math.top_k(loss, k=top_k_num)
    normalizer = tf.to_float(top_k_num)
    if num_valid_values is not None:
      normalizer = tf.minimum(normalizer, tf.to_float(num_valid_values))
    return tf.reduce_sum(top_k_losses) / (normalizer + 1e-5)
//  [python]
def aggregate_loss_top_k(self, loss, num_valid_values=None):
    """Aggregate the top-k the greatest loss values.

    Args:
      loss: a float tensor in shape (batch_size, num_boxes) or (batch_size,
        num_boxes, num_classes) which stores the loss values.
      num_valid_values: the number of loss values which are not ignored. The
        default value is None, which means all the loss values are valid.

    Returns:
      A 0-D float which stores the overall loss of the batch.
    """
    loss = tf.reshape(loss, shape=[-1])
    top_k_num = tf.cast(
        self._top_k_percent * tf.size(loss, out_type=tf.float32), tf.int32)
    top_k_losses, _ = tf.math.top_k(loss, k=top_k_num)
    normalizer = tf.cast(top_k_num, loss.dtype)
    if num_valid_values is not None:
      normalizer = tf.minimum(normalizer, tf.cast(num_valid_values, loss.dtype))
    return tf.reduce_sum(top_k_losses) / (normalizer + 1e-5)
//  [python]
def _fast_rcnn_box_loss(self, box_outputs, box_targets, class_targets,
                          normalizer=1.0):
    """Computes box regression loss."""
    with tf.name_scope('fast_rcnn_box_loss'):
      mask = tf.tile(
          tf.expand_dims(tf.greater(class_targets, 0), axis=2), [1, 1, 4])
      mask = tf.cast(mask, dtype=box_outputs.dtype)
      box_targets = tf.expand_dims(box_targets, axis=-1)
      box_outputs = tf.expand_dims(box_outputs, axis=-1)
      box_loss = self._huber_loss(box_targets, box_outputs, sample_weight=mask)
      # The loss is normalized by the number of ones in mask,
      # additional normalizer provided by the user and using 0.01 here to avoid
      # division by 0. For each replica, get the sum of non-zero masks. Then
      # get the mean of sums from all replicas. Note there is an extra division
      # by `num_replicas` in train_step(). So it is equivalent to normalizing
      # the box loss by the global sum of non-zero masks.
      replica_context = tf.distribute.get_replica_context()
      mask = tf.reduce_sum(mask)
      mask_mean = replica_context.all_reduce(
          tf.distribute.ReduceOp.MEAN, mask
      )
      box_loss /= normalizer * (mask_mean + 0.01)
      return box_loss
//  [python]
def normalize_boxes(boxes, image_shape):
  """Converts boxes to the normalized coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    normalized_boxes: a tensor whose shape is the same as `boxes` representing
      the normalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError(
        'boxes.shape[-1] is {:d}, but must be 4.'.format(boxes.shape[-1]))

  with tf.name_scope('normalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.to_float(image_shape)
      height = image_shape[..., 0:1]
      width = image_shape[..., 1:2]

    ymin = boxes[..., 0:1] / height
    xmin = boxes[..., 1:2] / width
    ymax = boxes[..., 2:3] / height
    xmax = boxes[..., 3:4] / width

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