jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def crf_forward(
    inputs,
    state,
    transition_params,
    sequence_lengths,
):
  """Computes the alpha values in a linear-chain CRF.

  See http://www.cs.columbia.edu/~mcollins/fb.pdf for reference.

  Args:
    inputs: A [batch_size, num_tags] matrix of unary potentials.
    state: A [batch_size, num_tags] matrix containing the previous alpha
      values.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.
      This matrix is expanded into a [1, num_tags, num_tags] in preparation
      for the broadcast summation occurring within the cell.
    sequence_lengths: A [batch_size] vector of true sequence lengths.

  Returns:
    new_alphas: A [batch_size, num_tags] matrix containing the
        new alpha values.
  """
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

  last_index = tf.maximum(
      tf.constant(0, dtype=sequence_lengths.dtype), sequence_lengths - 1)
  inputs = tf.transpose(inputs, [1, 0, 2])
  transition_params = tf.expand_dims(transition_params, 0)

  def _scan_fn(scan_state, scan_inputs):
    scan_state = tf.expand_dims(scan_state, 2)
    transition_scores = scan_state + transition_params
    new_alphas = scan_inputs + tf.reduce_logsumexp(transition_scores, [1])
    return new_alphas

  all_alphas = tf.transpose(tf.scan(_scan_fn, inputs, state), [1, 0, 2])
  # add first state for sequences of length 1
  all_alphas = tf.concat([tf.expand_dims(state, 1), all_alphas], 1)

  idxs = tf.stack([tf.range(tf.shape(last_index)[0]), last_index], axis=1)
  return tf.gather_nd(all_alphas, idxs)
//  [python]
def crf_decode_forward(
    inputs,
    state,
    transition_params,
    sequence_lengths,
):
  """Computes forward decoding in a linear-chain CRF.

  Args:
    inputs: A [batch_size, num_tags] matrix of unary potentials.
    state: A [batch_size, num_tags] matrix containing the previous step's
      score values.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.
    sequence_lengths: A [batch_size] vector of true sequence lengths.

  Returns:
    backpointers: A [batch_size, num_tags] matrix of backpointers.
    new_state: A [batch_size, num_tags] matrix of new score values.
  """
  sequence_lengths = tf.to_int32(sequence_lengths)
  mask = tf.sequence_mask(sequence_lengths, tf.shape(inputs)[1])
  crf_fwd_cell = CrfDecodeForwardRnnCell(transition_params, dtype=inputs.dtype)
  crf_fwd_layer = tf.keras.layers.RNN(
      crf_fwd_cell,
      return_sequences=True,
      return_state=True,
      dtype=inputs.dtype)
  return crf_fwd_layer(inputs, state, mask=mask)
//  [python]
def crf_decode_forward(
    inputs,
    state,
    transition_params,
    sequence_lengths,
):
  """Computes forward decoding in a linear-chain CRF.

  Args:
    inputs: A [batch_size, num_tags] matrix of unary potentials.
    state: A [batch_size, num_tags] matrix containing the previous step's
      score values.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.
    sequence_lengths: A [batch_size] vector of true sequence lengths.

  Returns:
    backpointers: A [batch_size, num_tags] matrix of backpointers.
    new_state: A [batch_size, num_tags] matrix of new score values.
  """
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
  mask = tf.sequence_mask(sequence_lengths, tf.shape(inputs)[1])
  crf_fwd_cell = CrfDecodeForwardRnnCell(transition_params, dtype=inputs.dtype)
  crf_fwd_layer = tf.keras.layers.RNN(
      crf_fwd_cell,
      return_sequences=True,
      return_state=True,
      dtype=inputs.dtype)
  return crf_fwd_layer(inputs, state, mask=mask)
//  [python]
def get_timestep_embedding(self, timestep, dim=320, max_period=10000):
        half = dim // 2
        log_max_period = tf.math.log(t.cast(max_period, tf.float32))
        freqs = tf.math.exp(
            -log_max_period * tf.range(0, half, dtype=tf.float32) / half
        )
        args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
        embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
        return embedding
//  [python]
def get_timestep_embedding(self, timestep, dim=320, max_period=10000):
        half = dim // 2
        log_max_period = tf.math.log(tf.cast(max_period, tf.float32))
        freqs = tf.math.exp(
            -log_max_period * tf.range(0, half, dtype=tf.float32) / half
        )
        args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
        embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
        return embedding
//  [python]
def compute_pairwise_distances(input_vecs: types.Tensor) -> types.Tensor:
  """Compute the pairwise distances matrix.

  Given input embedding vectors, this utility computes the (squared) pairwise
  distances matrix.

  Args:
    input_vecs: a `Tensor`. Input embedding vectors (one per row).

  Returns:
    The (squared) pairwise distances matrix. A dense float `Tensor` of shape
    [`num_vectors`, `num_vectors`], where `num_vectors` is the number of input
    embedding vectors.
  """
  r = tf.reduce_sum(input_vecs * input_vecs, axis=1, keepdims=True)
  pdistance_matrix = (
      r
      - 2 * tf.matmul(input_vecs, input_vecs, transpose_b=True)
      + tf.transpose(r)
  )
  return tf.to_float(pdistance_matrix)
//  [python]
def compute_pairwise_distances(input_vecs: types.Tensor) -> types.Tensor:
  """Compute the pairwise distances matrix.

  Given input embedding vectors, this utility computes the (squared) pairwise
  distances matrix.

  Args:
    input_vecs: a `Tensor`. Input embedding vectors (one per row).

  Returns:
    The (squared) pairwise distances matrix. A dense float `Tensor` of shape
    [`num_vectors`, `num_vectors`], where `num_vectors` is the number of input
    embedding vectors.
  """
  r = tf.reduce_sum(input_vecs * input_vecs, axis=1, keepdims=True)
  pdistance_matrix = (
      r
      - 2 * tf.matmul(input_vecs, input_vecs, transpose_b=True)
      + tf.transpose(r)
  )
  return tf.cast(pdistance_matrix, dtype=tf.float32)
//  [python]
def _test_forward_rel_op(data, func):
    with tf.Graph().as_default():
        in1 = tf.placeholder(shape=data[0].shape, dtype=data[0].dtype, name="in1")
        in2 = tf.placeholder(shape=data[1].shape, dtype=data[1].dtype, name="in2")
        op = func(in1, in2, name="op")
        _ = tf.to_int32(op, name="out1")
        compare_tf_with_tvm([data[0], data[1]], ["in1:0", "in2:0"], "out1:0")
//  [python]
def _test_forward_rel_op(data, func):
    with tf.Graph().as_default():
        in1 = tf.placeholder(shape=data[0].shape, dtype=data[0].dtype, name="in1")
        in2 = tf.placeholder(shape=data[1].shape, dtype=data[1].dtype, name="in2")
        op = func(in1, in2, name="op")
        _ = tf.cast(op, tf.int32, name="out1")
        compare_tf_with_tvm([data[0], data[1]], ["in1:0", "in2:0"], "out1:0")
//  [python]
def crf_decode(
    potentials,
    transition_params,
    sequence_length,
):
  """Decode the highest scoring sequence of tags.

  Args:
    potentials: A [batch_size, max_seq_len, num_tags] tensor of unary
      potentials.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.
    sequence_length: A [batch_size] vector of true sequence lengths.

  Returns:
    decode_tags: A [batch_size, max_seq_len] matrix, with dtype `tf.int32`.
                Contains the highest scoring tag indices.
    best_score: A [batch_size] vector, containing the score of `decode_tags`.
  """
  warnings.warn(
      "CRF Decoding does not work with KerasTensors in TF2.4. The bug has since been fixed in tensorflow/tensorflow##45534"
  )
  sequence_length = t.cast(sequence_length, dtype=t.int32)

  def _multi_seq_fn():
    # Computes forward decoding. Get last score and backpointers.
    initial_state = tf.slice(potentials, [0, 0, 0], [-1, 1, -1])
    initial_state = tf.squeeze(initial_state, axis=[1])
    inputs = tf.slice(potentials, [0, 1, 0], [-1, -1, -1])

    sequence_length_less_one = tf.maximum(
        tf.constant(0, dtype=tf.int32), sequence_length - 1)

    backpointers, last_score = crf_decode_forward(inputs, initial_state,
                                                  transition_params,
                                                  sequence_length_less_one)

    backpointers = tf.reverse_sequence(
        backpointers, sequence_length_less_one, seq_axis=1)

    initial_state = tf.cast(tf.argmax(last_score, axis=1), dtype=tf.int32)
    initial_state = tf.expand_dims(initial_state, axis=-1)

    decode_tags = crf_decode_backward(backpointers, initial_state)
    decode_tags = tf.squeeze(decode_tags, axis=[2])
    decode_tags = tf.concat([initial_state, decode_tags], axis=1)
    decode_tags = tf.reverse_sequence(decode_tags, sequence_length, seq_axis=1)

    best_score = tf.reduce_max(last_score, axis=1)
    return decode_tags, best_score

  return _multi_seq_fn()
//  [python]
def crf_decode(
    potentials,
    transition_params,
    sequence_length,
):
  """Decode the highest scoring sequence of tags.

  Args:
    potentials: A [batch_size, max_seq_len, num_tags] tensor of unary
      potentials.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.
    sequence_length: A [batch_size] vector of true sequence lengths.

  Returns:
    decode_tags: A [batch_size, max_seq_len] matrix, with dtype `tf.int32`.
                Contains the highest scoring tag indices.
    best_score: A [batch_size] vector, containing the score of `decode_tags`.
  """
  warnings.warn(
      "CRF Decoding does not work with KerasTensors in TF2.4. The bug has since been fixed in tensorflow/tensorflow##45534"
  )
  sequence_length = tf.cast(sequence_length, dtype=tf.int32)

  def _multi_seq_fn():
    # Computes forward decoding. Get last score and backpointers.
    initial_state = tf.slice(potentials, [0, 0, 0], [-1, 1, -1])
    initial_state = tf.squeeze(initial_state, axis=[1])
    inputs = tf.slice(potentials, [0, 1, 0], [-1, -1, -1])

    sequence_length_less_one = tf.maximum(
        tf.constant(0, dtype=tf.int32), sequence_length - 1)

    backpointers, last_score = crf_decode_forward(inputs, initial_state,
                                                  transition_params,
                                                  sequence_length_less_one)

    backpointers = tf.reverse_sequence(
        backpointers, sequence_length_less_one, seq_axis=1)

    initial_state = tf.cast(tf.argmax(last_score, axis=1), dtype=tf.int32)
    initial_state = tf.expand_dims(initial_state, axis=-1)

    decode_tags = crf_decode_backward(backpointers, initial_state)
    decode_tags = tf.squeeze(decode_tags, axis=[2])
    decode_tags = tf.concat([initial_state, decode_tags], axis=1)
    decode_tags = tf.reverse_sequence(decode_tags, sequence_length, seq_axis=1)

    best_score = tf.reduce_max(last_score, axis=1)
    return decode_tags, best_score

  return _multi_seq_fn()
//  [python]
def train_and_test(dataset_url, num_epochs, batch_size, evaluation_interval):
  """
  Train a model for training iterations with a batch size batch_size, printing accuracy every log_interval.
  :param dataset_url: The MNIST dataset url.
  :param num_epochs: The number of epochs to train for.
  :param batch_size: The batch size for training.
  :param evaluation_interval: The interval used to print the accuracy.
  :return:
  """

  with make_reader(os.path.join(dataset_url, 'train'), num_epochs=num_epochs) as train_reader:
    with make_reader(os.path.join(dataset_url, 'test'), num_epochs=num_epochs) as test_reader:
      # Create the model
      x = tf.placeholder(tf.float32, [None, 784])
      w = tf.Variable(tf.zeros([784, 10]))
      b = tf.Variable(tf.zeros([10]))
      y = tf.matmul(x, w) + b

      # Define loss and optimizer
      y_ = tf.placeholder(tf.int64, [None])

      # Define the loss function
      cross_entropy = tf.losses.sparse_softmax_cross_entropy(labels=y_, logits=y)

      train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)

      correct_prediction = tf.equal(tf.argmax(y, 1), y_)

      accuracy = tf.reduce_mean(tf.to_float(correct_prediction))

      train_dataset = make_dataset(train_reader) \
        .apply(tf.data.experimental.unbatch()) \
        .batch(batch_size) \
        .map(decode)

      train_iterator = train_dataset.make_one_shot_iterator()
      label, image = train_iterator.get_next()

      test_dataset = make_dataset(test_reader) \
        .apply(tf.data.experimental.unbatch()) \
        .batch(batch_size) \
        .map(decode)

      test_iterator = test_dataset.make_one_shot_iterator()
      test_label, test_image = test_iterator.get_next()

      # Train
      print('Training model for {0} epoch with batch size {1} and evaluation interval {2}'.format(
        num_epochs, batch_size, evaluation_interval
      ))

      i = 0
      with tf.Session() as sess:
        sess.run([
          tf.local_variables_initializer(),
          tf.global_variables_initializer(),
        ])

        try:
          while True:
            cur_label, cur_image = sess.run([label, image])

            sess.run([train_step], feed_dict={x: cur_image, y_:cur_label})

            if i % evaluation_interval == 0:
              test_cur_label, test_cur_image = sess.run([test_label, test_image])
              print('After {0} training iterations, the accuracy of the model is: {1:.2f}'.format(
                i,
                sess.run(accuracy, feed_dict={
                  x: test_cur_image, y_: test_cur_label})))
            i += 1

        except tf.errors.OutOfRangeError:
          print("Finish! the number is " + str(i))
//  [python]
def train_and_test(dataset_url, num_epochs, batch_size, evaluation_interval):
  """
  Train a model for training iterations with a batch size batch_size, printing accuracy every log_interval.
  :param dataset_url: The MNIST dataset url.
  :param num_epochs: The number of epochs to train for.
  :param batch_size: The batch size for training.
  :param evaluation_interval: The interval used to print the accuracy.
  :return:
  """

  with make_reader(os.path.join(dataset_url, 'train'), num_epochs=num_epochs) as train_reader:
    with make_reader(os.path.join(dataset_url, 'test'), num_epochs=num_epochs) as test_reader:
      # Create the model
      x = tf.placeholder(tf.float32, [None, 784])
      w = tf.Variable(tf.zeros([784, 10]))
      b = tf.Variable(tf.zeros([10]))
      y = tf.matmul(x, w) + b

      # Define loss and optimizer
      y_ = tf.placeholder(tf.int64, [None])

      # Define the loss function
      cross_entropy = tf.losses.sparse_softmax_cross_entropy(labels=y_, logits=y)

      train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)

      correct_prediction = tf.equal(tf.argmax(y, 1), y_)

      accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

      train_dataset = make_dataset(train_reader) \
        .apply(tf.data.experimental.unbatch()) \
        .batch(batch_size) \
        .map(decode)

      train_iterator = train_dataset.make_one_shot_iterator()
      label, image = train_iterator.get_next()

      test_dataset = make_dataset(test_reader) \
        .apply(tf.data.experimental.unbatch()) \
        .batch(batch_size) \
        .map(decode)

      test_iterator = test_dataset.make_one_shot_iterator()
      test_label, test_image = test_iterator.get_next()

      # Train
      print('Training model for {0} epoch with batch size {1} and evaluation interval {2}'.format(
        num_epochs, batch_size, evaluation_interval
      ))

      i = 0
      with tf.Session() as sess:
        sess.run([
          tf.local_variables_initializer(),
          tf.global_variables_initializer(),
        ])

        try:
          while True:
            cur_label, cur_image = sess.run([label, image])

            sess.run([train_step], feed_dict={x: cur_image, y_:cur_label})

            if i % evaluation_interval == 0:
              test_cur_label, test_cur_image = sess.run([test_label, test_image])
              print('After {0} training iterations, the accuracy of the model is: {1:.2f}'.format(
                i,
                sess.run(accuracy, feed_dict={
                  x: test_cur_image, y_: test_cur_label})))
            i += 1

        except tf.errors.OutOfRangeError:
          print("Finish! the number is " + str(i))
//  [python]
def _shift_right(self, input_ids):
        decoder_start_token_id = self.decoder_start_token_id
        pad_token_id = self.decoder_pad_token_id

        if decoder_start_token_id is None or pad_token_id is None:
            raise ValueError("decoder_start_token_id and pad_token_id must be defined!")

        start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)
        start_tokens = tf.to_int32(start_tokens)  # Ensure compatible dtypes for concatenation
        shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)

        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100,
            tf.to_int32(tf.fill(shape_list(shifted_input_ids), pad_token_id)),
            shifted_input_ids,
        )

        # "Verify that `labels` has only positive values and -100"
        tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=shifted_input_ids.dtype))

        return shifted_input_ids
//  [python]
def _shift_right(self, input_ids):
        decoder_start_token_id = self.decoder_start_token_id
        pad_token_id = self.decoder_pad_token_id

        if decoder_start_token_id is None or pad_token_id is None:
            raise ValueError("decoder_start_token_id and pad_token_id must be defined!")

        start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)
        start_tokens = tf.cast(start_tokens, input_ids.dtype)  # Ensure compatible dtypes for concatenation
        shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)

        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100,
            tf.cast(tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids.dtype),
            shifted_input_ids,
        )

        # "Verify that `labels` has only positive values and -100"
        tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=shifted_input_ids.dtype))

        return shifted_input_ids
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_mask': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'segment_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
    }
    if self._xlnet_format:
      name_to_features['class_index'] = tf.io.FixedLenFeature([], tf.int64)
      name_to_features['paragraph_mask'] = tf.io.FixedLenFeature(
          [self._seq_length], tf.int64)
      if self._is_training:
        name_to_features['is_impossible'] = tf.io.FixedLenFeature([], tf.int64)

    if self._is_training:
      name_to_features['start_positions'] = tf.io.FixedLenFeature([], tf.int64)
      name_to_features['end_positions'] = tf.io.FixedLenFeature([], tf.int64)
    else:
      name_to_features['unique_ids'] = tf.io.FixedLenFeature([], tf.int64)
    example = tf.io.parse_single_example(record, name_to_features)

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

    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_mask': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'segment_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
    }
    if self._xlnet_format:
      name_to_features['class_index'] = tf.io.FixedLenFeature([], tf.int64)
      name_to_features['paragraph_mask'] = tf.io.FixedLenFeature(
          [self._seq_length], tf.int64)
      if self._is_training:
        name_to_features['is_impossible'] = tf.io.FixedLenFeature([], tf.int64)

    if self._is_training:
      name_to_features['start_positions'] = tf.io.FixedLenFeature([], tf.int64)
      name_to_features['end_positions'] = tf.io.FixedLenFeature([], tf.int64)
    else:
      name_to_features['unique_ids'] = tf.io.FixedLenFeature([], tf.int64)
    example = tf.io.parse_single_example(record, name_to_features)

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

    return example
//  [python]
def ms_loss(labels, embeddings, alpha=2.0, beta=50.0, lamb=1.0, eps=0.1, ms_mining=False, embed_normed=True):
    """
    ref: http://openaccess.thecvf.com/content_CVPR_2019/papers/Wang_Multi-Similarity_Loss_With_General_Pair_Weighting_for_Deep_Metric_Learning_CVPR_2019_paper.pdf
    official codes: https://github.com/MalongTech/research-ms-loss
    """
    # make sure embedding should be l2-normalized
    if not embed_normed:
        embeddings = tf.nn.l2_normalize(embeddings, axis=1)
    labels = tf.reshape(labels, [-1, 1])

    embed_shape = modeling.get_shape_list(embeddings)
    batch_size = embed_shape[0]

    adjacency = tf.equal(labels, tf.transpose(labels))
    adjacency_not = tf.logical_not(adjacency)

    mask_pos = tf.to_float(adjacency) - tf.eye(batch_size, dtype=tf.float32)
    mask_neg = tf.to_float(adjacency_not)

    sim_mat = tf.matmul(embeddings, embeddings, transpose_a=False, transpose_b=True)
    sim_mat = tf.maximum(sim_mat, 0.0)

    pos_mat = tf.multiply(sim_mat, mask_pos)
    neg_mat = tf.multiply(sim_mat, mask_neg)

    if ms_mining:
        max_val = tf.reduce_max(neg_mat, axis=1, keepdims=True)
        tmp_max_val = tf.reduce_max(pos_mat, axis=1, keepdims=True)
        min_val = tf.reduce_min(tf.multiply(sim_mat - tmp_max_val, mask_pos), axis=1, keepdims=True) + tmp_max_val

        max_val = tf.tile(max_val, [1, batch_size])
        min_val = tf.tile(min_val, [1, batch_size])

        mask_pos = tf.where(pos_mat < max_val + eps, mask_pos, tf.zeros_like(mask_pos))
        mask_neg = tf.where(neg_mat > min_val - eps, mask_neg, tf.zeros_like(mask_neg))

    pos_exp = tf.exp(-alpha * (pos_mat - lamb))
    pos_exp = tf.where(mask_pos > 0.0, pos_exp, tf.zeros_like(pos_exp))

    neg_exp = tf.exp(beta * (neg_mat - lamb))
    neg_exp = tf.where(mask_neg > 0.0, neg_exp, tf.zeros_like(neg_exp))

    pos_term = tf.log(1.0 + tf.reduce_sum(pos_exp, axis=1)) / alpha
    neg_term = tf.log(1.0 + tf.reduce_sum(neg_exp, axis=1)) / beta

    loss = tf.reduce_mean(pos_term + neg_term)
    return loss
//  [python]
def ms_loss(labels, embeddings, alpha=2.0, beta=50.0, lamb=1.0, eps=0.1, ms_mining=False, embed_normed=True):
    """
    ref: http://openaccess.thecvf.com/content_CVPR_2019/papers/Wang_Multi-Similarity_Loss_With_General_Pair_Weighting_for_Deep_Metric_Learning_CVPR_2019_paper.pdf
    official codes: https://github.com/MalongTech/research-ms-loss
    """
    # make sure embedding should be l2-normalized
    if not embed_normed:
        embeddings = tf.nn.l2_normalize(embeddings, axis=1)
    labels = tf.reshape(labels, [-1, 1])

    embed_shape = modeling.get_shape_list(embeddings)
    batch_size = embed_shape[0]

    adjacency = tf.equal(labels, tf.transpose(labels))
    adjacency_not = tf.logical_not(adjacency)

    mask_pos = tf.cast(adjacency, dtype=tf.float32) - tf.eye(batch_size, dtype=tf.float32)
    mask_neg = tf.cast(adjacency_not, dtype=tf.float32)

    sim_mat = tf.matmul(embeddings, embeddings, transpose_a=False, transpose_b=True)
    sim_mat = tf.maximum(sim_mat, 0.0)

    pos_mat = tf.multiply(sim_mat, mask_pos)
    neg_mat = tf.multiply(sim_mat, mask_neg)

    if ms_mining:
        max_val = tf.reduce_max(neg_mat, axis=1, keepdims=True)
        tmp_max_val = tf.reduce_max(pos_mat, axis=1, keepdims=True)
        min_val = tf.reduce_min(tf.multiply(sim_mat - tmp_max_val, mask_pos), axis=1, keepdims=True) + tmp_max_val

        max_val = tf.tile(max_val, [1, batch_size])
        min_val = tf.tile(min_val, [1, batch_size])

        mask_pos = tf.where(pos_mat < max_val + eps, mask_pos, tf.zeros_like(mask_pos))
        mask_neg = tf.where(neg_mat > min_val - eps, mask_neg, tf.zeros_like(mask_neg))

    pos_exp = tf.exp(-alpha * (pos_mat - lamb))
    pos_exp = tf.where(mask_pos > 0.0, pos_exp, tf.zeros_like(pos_exp))

    neg_exp = tf.exp(beta * (neg_mat - lamb))
    neg_exp = tf.where(mask_neg > 0.0, neg_exp, tf.zeros_like(neg_exp))

    pos_term = tf.log(1.0 + tf.reduce_sum(pos_exp, axis=1)) / alpha
    neg_term = tf.log(1.0 + tf.reduce_sum(neg_exp, axis=1)) / beta

    loss = tf.reduce_mean(pos_term + neg_term)
    return loss
//  [python]
def _learning_rate_scheduler(self, global_step):
        d_model = t.cast(self.d_model, t.float32)
        step = t.cast(global_step, t.float32)
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warm_up_steps ** -1.5)
        return tf.math.rsqrt(d_model) * tf.math.minimum(arg1, arg2)
//  [python]
def _learning_rate_scheduler(self, global_step):
        d_model = tf.cast(self.d_model, tf.float32)
        step = tf.cast(global_step, tf.float32)
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warm_up_steps ** -1.5)
        return tf.math.rsqrt(d_model) * tf.math.minimum(arg1, arg2)
//  [python]
def _drop_attention_logits(
        self, logits: tf.Tensor, pad_mask: tf.Tensor, training: tf.Tensor
    ) -> tf.Tensor:
        def droped_logits() -> tf.Tensor:
            keep_prob = tf.random.uniform(tf.shape(logits), 0, 1) + pad_mask
            drop_mask = tf.to_float(
                tf.less(keep_prob, self.attention_dropout_rate)
            )

            return logits + drop_mask * -1e9

        return smart_cond(training, droped_logits, lambda: tf.identity(logits))
//  [python]
def _drop_attention_logits(
        self, logits: tf.Tensor, pad_mask: tf.Tensor, training: tf.Tensor
    ) -> tf.Tensor:
        def droped_logits() -> tf.Tensor:
            keep_prob = tf.random.uniform(tf.shape(logits), 0, 1) + pad_mask
            drop_mask = tf.cast(
                tf.less(keep_prob, self.attention_dropout_rate), logits.dtype
            )

            return logits + drop_mask * -1e9

        return smart_cond(training, droped_logits, lambda: tf.identity(logits))
//  [python]
def positional_encoding(self):
        '''
        calculate the positional encoding given the window length
        :return: positional encoding (1, window_length, 1, d_model)
        '''
        angle_rads = self.get_angles(np.arange(self.window_len)[:, np.newaxis], np.arange(self.d_model)[np.newaxis, :])

        # apply sin to even indices in the array; 2i
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

        pos_encoding = angle_rads[np.newaxis, :, np.newaxis, :]

        return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(self):
        '''
        calculate the positional encoding given the window length
        :return: positional encoding (1, window_length, 1, d_model)
        '''
        angle_rads = self.get_angles(np.arange(self.window_len)[:, np.newaxis], np.arange(self.d_model)[np.newaxis, :])

        # apply sin to even indices in the array; 2i
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

        pos_encoding = angle_rads[np.newaxis, :, np.newaxis, :]

        return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
        # attn_dim: num_joints for spatial and seq_len for temporal
        '''
        The scaled dot product attention mechanism introduced in the Transformer
        :param q: the query vectors matrix (..., attn_dim, d_model/num_heads)
        :param k: the key vector matrix (..., attn_dim, d_model/num_heads)
        :param v: the value vector matrix (..., attn_dim, d_model/num_heads)
        :param mask: a mask for attention
        :return: the updated encoding and the attention weights matrix
        '''

        matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., num_heads, attn_dim, attn_dim)

        # scale matmul_qk
        dk = tf.to_float(tf.shape(k)[-1])
        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

        # add the mask to the scaled tensor.
        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        # normalized on the last axis (seq_len_k) so that the scores add up to 1.
        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., num_heads, attn_dim, attn_dim)

        output = tf.matmul(attention_weights, v)  # (..., num_heads, attn_dim, depth)

        return output, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
        # attn_dim: num_joints for spatial and seq_len for temporal
        '''
        The scaled dot product attention mechanism introduced in the Transformer
        :param q: the query vectors matrix (..., attn_dim, d_model/num_heads)
        :param k: the key vector matrix (..., attn_dim, d_model/num_heads)
        :param v: the value vector matrix (..., attn_dim, d_model/num_heads)
        :param mask: a mask for attention
        :return: the updated encoding and the attention weights matrix
        '''

        matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., num_heads, attn_dim, attn_dim)

        # scale matmul_qk
        dk = tf.cast(tf.shape(k)[-1], tf.float32)
        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

        # add the mask to the scaled tensor.
        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        # normalized on the last axis (seq_len_k) so that the scores add up to 1.
        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., num_heads, attn_dim, attn_dim)

        output = tf.matmul(attention_weights, v)  # (..., num_heads, attn_dim, depth)

        return output, attention_weights
//  [python]
def _positional_encoding(self, max_position: tf.Tensor) -> tf.Tensor:
        max_position = t.cast(max_position, dtype=tf.float32)
        angle_rads = tf.range(max_position)[:, tf.newaxis] * self._angles

        # transpose for easy slicing
        angle_rads = tf.transpose(angle_rads, perm=[1, 0])
        shape = tf.shape(angle_rads)
        # apply sin to even indices in the array; 2i
        sin_even = tf.sin(tf.gather_nd(angle_rads, self._even_indices))
        pos_encoding_even = tf.scatter_nd(self._even_indices, sin_even, shape)
        # apply cos to odd indices in the array; 2i+1
        cos_odd = tf.cos(tf.gather_nd(angle_rads, self._odd_indices))
        pos_encoding_odd = tf.scatter_nd(self._odd_indices, cos_odd, shape)
        # combine even and odd positions and transpose back
        pos_encoding = tf.transpose(pos_encoding_even + pos_encoding_odd, perm=[1, 0])
        # add batch dimension
        return tf.stop_gradient(pos_encoding[tf.newaxis, ...])
//  [python]
def _positional_encoding(self, max_position: tf.Tensor) -> tf.Tensor:
        max_position = tf.cast(max_position, dtype=tf.float32)
        angle_rads = tf.range(max_position)[:, tf.newaxis] * self._angles

        # transpose for easy slicing
        angle_rads = tf.transpose(angle_rads, perm=[1, 0])
        shape = tf.shape(angle_rads)
        # apply sin to even indices in the array; 2i
        sin_even = tf.sin(tf.gather_nd(angle_rads, self._even_indices))
        pos_encoding_even = tf.scatter_nd(self._even_indices, sin_even, shape)
        # apply cos to odd indices in the array; 2i+1
        cos_odd = tf.cos(tf.gather_nd(angle_rads, self._odd_indices))
        pos_encoding_odd = tf.scatter_nd(self._odd_indices, cos_odd, shape)
        # combine even and odd positions and transpose back
        pos_encoding = tf.transpose(pos_encoding_even + pos_encoding_odd, perm=[1, 0])
        # add batch dimension
        return tf.stop_gradient(pos_encoding[tf.newaxis, ...])
//  [python]
def _scaled_dot_product_attention(
        self,
        query: tf.Tensor,
        key: tf.Tensor,
        value: tf.Tensor,
        pad_mask: tf.Tensor,
        training: tf.Tensor,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        """Calculate the attention weights.

        query, key, value must have matching leading dimensions.
        key, value must have matching penultimate dimension,
        i.e.: seq_len_k = seq_len_v.
        The mask has different shapes depending on its type (padding or look ahead)
        but it must be broadcastable for addition.

        Arguments:
            query: A tensor with shape (..., length, depth).
            key: A tensor with shape (..., length, depth).
            value: A tensor with shape (..., length, depth).
            pad_mask: Float tensor with shape broadcastable
                to (..., length, length). Defaults to None.

        Returns:
            output: A tensor with shape (..., length, depth).
            attention_weights: A tensor with shape (..., length, length).
        """
        matmul_qk = tf.matmul(query, key, transpose_b=True)  # (..., length, length)

        if self.use_key_relative_position:
            matmul_qk += self._matmul_with_relative_keys(query)

        # scale matmul_qk
        dk = tf.to_float(tf.shape(key)[-1])
        logits = matmul_qk / tf.math.sqrt(dk)

        # add the mask to the scaled tensor.
        if pad_mask is not None:
            logits += pad_mask * -1e9

        # apply attention dropout before softmax to maintain attention_weights norm as 1
        if self.attention_dropout_rate > 0:
            logits = self._drop_attention_logits(logits, pad_mask, training)

        # softmax is normalized on the last axis (length) so that the scores
        # add up to 1.
        attention_weights = tf.nn.softmax(logits, axis=-1)  # (..., length, length)

        output = tf.matmul(attention_weights, value)  # (..., length, depth)
        if self.use_value_relative_position:
            output += self._matmul_with_relative_values(attention_weights)

        return output, attention_weights
//  [python]
def _scaled_dot_product_attention(
        self,
        query: tf.Tensor,
        key: tf.Tensor,
        value: tf.Tensor,
        pad_mask: tf.Tensor,
        training: tf.Tensor,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        """Calculate the attention weights.

        query, key, value must have matching leading dimensions.
        key, value must have matching penultimate dimension,
        i.e.: seq_len_k = seq_len_v.
        The mask has different shapes depending on its type (padding or look ahead)
        but it must be broadcastable for addition.

        Arguments:
            query: A tensor with shape (..., length, depth).
            key: A tensor with shape (..., length, depth).
            value: A tensor with shape (..., length, depth).
            pad_mask: Float tensor with shape broadcastable
                to (..., length, length). Defaults to None.

        Returns:
            output: A tensor with shape (..., length, depth).
            attention_weights: A tensor with shape (..., length, length).
        """
        matmul_qk = tf.matmul(query, key, transpose_b=True)  # (..., length, length)

        if self.use_key_relative_position:
            matmul_qk += self._matmul_with_relative_keys(query)

        # scale matmul_qk
        dk = tf.cast(tf.shape(key)[-1], tf.float32)
        logits = matmul_qk / tf.math.sqrt(dk)

        # add the mask to the scaled tensor.
        if pad_mask is not None:
            logits += pad_mask * -1e9

        # apply attention dropout before softmax to maintain attention_weights norm as 1
        if self.attention_dropout_rate > 0:
            logits = self._drop_attention_logits(logits, pad_mask, training)

        # softmax is normalized on the last axis (length) so that the scores
        # add up to 1.
        attention_weights = tf.nn.softmax(logits, axis=-1)  # (..., length, length)

        output = tf.matmul(attention_weights, value)  # (..., length, depth)
        if self.use_value_relative_position:
            output += self._matmul_with_relative_values(attention_weights)

        return output, attention_weights
//  [python]
def causal_attention_mask(nd, ns, dtype):
        """
        1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),
        -1, ns-nd), but doesn't produce garbage on TPUs.
        """
        i = tf.range(nd)[:, None]
        j = tf.range(ns)
        m = i >= j - ns + nd
        return t.cast(m, dtype)
//  [python]
def causal_attention_mask(nd, ns, dtype):
        """
        1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),
        -1, ns-nd), but doesn't produce garbage on TPUs.
        """
        i = tf.range(nd)[:, None]
        j = tf.range(ns)
        m = i >= j - ns + nd
        return tf.cast(m, dtype)
//  [python]
def call(
        self,
        x: tf.Tensor,
        pad_mask: Optional[tf.Tensor] = None,
        training: Optional[Union[tf.Tensor, bool]] = None,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        """Apply transformer encoder.

        Arguments:
            x: A tensor with shape [batch_size, length, input_size].
            pad_mask: Float tensor with shape broadcastable
                to (..., length, length). Defaults to None.
            training: A bool, whether in training mode or not.

        Returns:
            Transformer encoder output with shape [batch_size, length, units]
        """
        # adding embedding and position encoding.
        x = self._embedding(x)  # (batch_size, length, units)
        x *= tf.sqrt(tf.to_float(self.units))
        x += self._positional_encoding(tf.shape(x)[1])
        x = self._dropout(x, training=training)

        if pad_mask is not None:
            pad_mask = tf.squeeze(pad_mask, -1)  # (batch_size, length)
            pad_mask = pad_mask[:, tf.newaxis, tf.newaxis, :]
            # pad_mask.shape = (batch_size, 1, 1, length)
            if self.unidirectional:
                # add look ahead pad mask to emulate unidirectional behavior
                pad_mask = tf.minimum(
                    1.0, pad_mask + self._look_ahead_pad_mask(tf.shape(pad_mask)[-1])
                )  # (batch_size, 1, length, length)

        layer_attention_weights = []

        for layer in self._enc_layers:
            x, attn_weights = layer(x, pad_mask=pad_mask, training=training)
            layer_attention_weights.append(attn_weights)

        # if normalization is done in encoding layers, then it should also be done
        # on the output, since the output can grow very large, being the sum of
        # a whole stack of unnormalized layer outputs.
        x = self._layer_norm(x)  # (batch_size, length, units)

        # Keep the batch dimension on the first axis
        attention_weights_as_output = tf.transpose(
            tf.stack(layer_attention_weights), (1, 0, 2, 3, 4)
        )

        # (batch_size, length, units),
        # (batch_size, num_layers, num_heads, length, length)
        return x, attention_weights_as_output
//  [python]
def call(
        self,
        x: tf.Tensor,
        pad_mask: Optional[tf.Tensor] = None,
        training: Optional[Union[tf.Tensor, bool]] = None,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        """Apply transformer encoder.

        Arguments:
            x: A tensor with shape [batch_size, length, input_size].
            pad_mask: Float tensor with shape broadcastable
                to (..., length, length). Defaults to None.
            training: A bool, whether in training mode or not.

        Returns:
            Transformer encoder output with shape [batch_size, length, units]
        """
        # adding embedding and position encoding.
        x = self._embedding(x)  # (batch_size, length, units)
        x *= tf.math.sqrt(tf.cast(self.units, tf.float32))
        x += self._positional_encoding(tf.shape(x)[1])
        x = self._dropout(x, training=training)

        if pad_mask is not None:
            pad_mask = tf.squeeze(pad_mask, -1)  # (batch_size, length)
            pad_mask = pad_mask[:, tf.newaxis, tf.newaxis, :]
            # pad_mask.shape = (batch_size, 1, 1, length)
            if self.unidirectional:
                # add look ahead pad mask to emulate unidirectional behavior
                pad_mask = tf.minimum(
                    1.0, pad_mask + self._look_ahead_pad_mask(tf.shape(pad_mask)[-1])
                )  # (batch_size, 1, length, length)

        layer_attention_weights = []

        for layer in self._enc_layers:
            x, attn_weights = layer(x, pad_mask=pad_mask, training=training)
            layer_attention_weights.append(attn_weights)

        # if normalization is done in encoding layers, then it should also be done
        # on the output, since the output can grow very large, being the sum of
        # a whole stack of unnormalized layer outputs.
        x = self._layer_norm(x)  # (batch_size, length, units)

        # Keep the batch dimension on the first axis
        attention_weights_as_output = tf.transpose(
            tf.stack(layer_attention_weights), (1, 0, 2, 3, 4)
        )

        # (batch_size, length, units),
        # (batch_size, num_layers, num_heads, length, length)
        return x, attention_weights_as_output
//  [python]
def get_timestep_embedding(self, timestep, dim=320, max_period=10000):
    half = dim // 2
    log_max_preiod = tf.math.log(tf.cast(max_period, tf.float32))
    # The docker builds could not support unary `-`.
    # pylint: disable=invalid-unary-operand-type
    freqs = tf.math.exp(
        -log_max_preiod * tf.range(0, half, dtype=tf.float32) / half
    )
    args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
    embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
    embedding = tf.reshape(embedding, [1, -1])
    return embedding
//  [python]
def _get_add_noise(stddev, seed: Optional[int] = None):
  """Utility function to decide which `add_noise` to use according to tf version."""
  if distutils.version.LooseVersion(
      tf.__version__) < distutils.version.LooseVersion('2.0.0'):

    # The seed should be only used for testing purpose.
    if seed is not None:
      tf.random.set_seed(seed)

    def add_noise(v):
      return v + tf.random.normal(
          tf.shape(input=v), stddev=stddev, dtype=v.dtype)
  else:
    random_normal = tf.random_normal_initializer(stddev=stddev, seed=seed)

    def add_noise(v):
      return v + tf.to_float(random_normal(tf.shape(input=v)))

  return add_noise
//  [python]
def _get_add_noise(stddev, seed: Optional[int] = None):
  """Utility function to decide which `add_noise` to use according to tf version."""
  if distutils.version.LooseVersion(
      tf.__version__) < distutils.version.LooseVersion('2.0.0'):

    # The seed should be only used for testing purpose.
    if seed is not None:
      tf.random.set_seed(seed)

    def add_noise(v):
      return v + tf.random.normal(
          tf.shape(input=v), stddev=stddev, dtype=v.dtype)
  else:
    random_normal = tf.random_normal_initializer(stddev=stddev, seed=seed)

    def add_noise(v):
      return v + tf.cast(random_normal(tf.shape(input=v)), dtype=v.dtype)

  return add_noise
//  [python]
def per_instance_masks(panoptic_labels: tf.Tensor,
                       instance_panoptic_labels: tf.Tensor,
                       out_dtype: tf.dtypes.DType = tf.bool) -> tf.Tensor:
  """3D tensor where each slice in 3rd dimensions is an instance's mask."""
  num_instances = tf.size(instance_panoptic_labels)
  matches = tf.equal(
      tf.expand_dims(panoptic_labels, 0),
      tf.reshape(instance_panoptic_labels, [num_instances, 1, 1]))
  return tf.to_int64(matches)
//  [python]
def per_instance_masks(panoptic_labels: tf.Tensor,
                       instance_panoptic_labels: tf.Tensor,
                       out_dtype: tf.dtypes.DType = tf.bool) -> tf.Tensor:
  """3D tensor where each slice in 3rd dimensions is an instance's mask."""
  num_instances = tf.size(instance_panoptic_labels)
  matches = tf.equal(
      tf.expand_dims(panoptic_labels, 0),
      tf.reshape(instance_panoptic_labels, [num_instances, 1, 1]))
  return tf.cast(matches, out_dtype)
//  [python]
def _attn(self, q, k, v, attention_mask, head_mask, output_attentions, training=False):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        if self.scale:
            dk = tf.cast(shape_list(k)[-1], dtype=w.dtype)  # scale attention_scores
            w = w / tf.math.sqrt(dk)

        if not self.is_cross_attention:
            # if only "normal" attention layer implements causal mask

            # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
            _, _, nd, ns = shape_list(w)
            b = self.causal_attention_mask(nd, ns, dtype=w.dtype)
            b = tf.reshape(b, [1, 1, nd, ns])
            w = w * b - 1e4 * (1 - b)

        if attention_mask is not None:
            # Apply the attention mask
            attention_mask = tf.cast(attention_mask, dtype=w.dtype)
            w = w + attention_mask

        w = stable_softmax(w, axis=-1)
        w = self.attn_dropout(w, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            w = w * head_mask

        outputs = [tf.matmul(w, v)]
        if output_attentions:
            outputs.append(w)
        return outputs
//  [python]
def sample(self, amount=None):
        amount = [amount] if amount else []
        indices = self._dist.sample(*amount)
        sample = self._one_hot(indices)
        probs = self._dist.probs_parameter()
        sample += t.cast(probs - tf.stop_gradient(probs), self._dtype)
        return sample
//  [python]
def sample(self, amount=None):
        amount = [amount] if amount else []
        indices = self._dist.sample(*amount)
        sample = self._one_hot(indices)
        probs = self._dist.probs_parameter()
        sample += tf.cast(probs - tf.stop_gradient(probs), self._dtype)
        return sample
//  [python]
def _inverse(self, y):
        dtype = y.dtype
        y = t.cast(y, tf.float32)
        y = tf.where(tf.less_equal(tf.abs(y), 1.0), tf.clip_by_value(y, -0.99999997, 0.99999997), y)
        y = tf.atanh(y)
        y = t.cast(y, dtype)
        return y
//  [python]
def _inverse(self, y):
        dtype = y.dtype
        y = tf.cast(y, tf.float32)
        y = tf.where(tf.less_equal(tf.abs(y), 1.0), tf.clip_by_value(y, -0.99999997, 0.99999997), y)
        y = tf.atanh(y)
        y = tf.cast(y, dtype)
        return y
//  [python]
def build(self, input_shape: tf.TensorShape):
        """
        Build shared token embedding layer Shared weights logic adapted from
        https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24
        """

        weight = self._init_weight(self.num_positions, self.embedding_dim)

        self.weight = self.add_weight(
            name="embeddings",
            shape=[self.num_positions, self.embedding_dim],
        )
        weight = tf.to_float(weight)

        self.weight.assign(weight)

        super().build(input_shape)
//  [python]
def build(self, input_shape: tf.TensorShape):
        """
        Build shared token embedding layer Shared weights logic adapted from
        https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24
        """

        weight = self._init_weight(self.num_positions, self.embedding_dim)

        self.weight = self.add_weight(
            name="embeddings",
            shape=[self.num_positions, self.embedding_dim],
        )
        weight = tf.cast(weight, dtype=self.weight.dtype)

        self.weight.assign(weight)

        super().build(input_shape)