jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def crop_normalized_boxes(
    bboxes: tf.Tensor,
    ori_image_size: tf.Tensor,
    new_image_size: tf.Tensor,
    offset: tf.Tensor,
) -> tf.Tensor:
  """Crop normalized boxes.

  Args:
    bboxes: a float tensor of shape [N, 4] representing normalized box
      coordinates.
    ori_image_size: an int tensor of shape [2] representing the original image
      size.
    new_image_size: an int tensor of shape [2] representing the cropped image
      size.
    offset: an int tensor of shape [2] representing the offset of the crop.

  Returns:
    a tensor of shape [N, 4] representing the new normalized bounding box
    coordinates in the new cropped image.
  """
  new_bboxes = box_ops.denormalize_boxes(bboxes, ori_image_size)
  new_bboxes -= tf.tile(tf.cast(offset, dtype=tf.float32), [2])[None, ...]
  new_bboxes = box_ops.normalize_boxes(new_bboxes, new_image_size)
  return tf.clip_by_value(new_bboxes, 0.0, 1.0)
//  [python]
def test_attention_projection(
      self, feature_transform, num_random_features, training, redraw, is_short,
      begin_kernel):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=redraw,
        is_short_seq=is_short,
        begin_kernel=begin_kernel)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    masks = tf.to_float(encoder_inputs_mask)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def test_attention_projection(
      self, feature_transform, num_random_features, training, redraw, is_short,
      begin_kernel):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=redraw,
        is_short_seq=is_short,
        begin_kernel=begin_kernel)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    masks = tf.cast(encoder_inputs_mask, dtype=tf.float32)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def argsort(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: int = -1,
    descending: bool = False,
    stable: bool = True,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    direction = "DESCENDING" if descending else "ASCENDING"
    x = tf.convert_to_tensor(x)
    ret = tf.argsort(x, axis=axis, direction=direction, stable=stable)
    return tf.cast(ret, dtype=tf.int64)
//  [python]
def build_dataset(self, filenames, labels, is_training):
    """Build input dataset."""
    filenames = tf.constant(filenames)
    labels = tf.constant(labels)
    dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))

    def _parse_function(filename, label):
      image_string = tf.read_file(filename)
      image_decoded = preprocessing.preprocess_image(
          image_string, is_training, self.image_size)
      image = tf.to_float(image_decoded)
      return image, label

    dataset = dataset.map(_parse_function)
    dataset = dataset.batch(self.batch_size)

    iterator = dataset.make_one_shot_iterator()
    images, labels = iterator.get_next()
    return images, labels
//  [python]
def build_dataset(self, filenames, labels, is_training):
    """Build input dataset."""
    filenames = tf.constant(filenames)
    labels = tf.constant(labels)
    dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))

    def _parse_function(filename, label):
      image_string = tf.read_file(filename)
      image_decoded = preprocessing.preprocess_image(
          image_string, is_training, self.image_size)
      image = tf.cast(image_decoded, tf.float32)
      return image, label

    dataset = dataset.map(_parse_function)
    dataset = dataset.batch(self.batch_size)

    iterator = dataset.make_one_shot_iterator()
    images, labels = iterator.get_next()
    return images, labels
//  [python]
def test_causal_windowed_attention_projection(
      self, feature_transform, num_random_features, training, redraw,
      begin_kernel, causal_window_decay, causal_padding):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=redraw,
        is_short_seq=False,
        begin_kernel=begin_kernel,
        use_causal_windowed=True,
        causal_chunk_length=8,
        causal_window_length=3,
        causal_window_decay=causal_window_decay,
        causal_padding=causal_padding)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    masks = tf.to_float(encoder_inputs_mask)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def test_causal_windowed_attention_projection(
      self, feature_transform, num_random_features, training, redraw,
      begin_kernel, causal_window_decay, causal_padding):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=redraw,
        is_short_seq=False,
        begin_kernel=begin_kernel,
        use_causal_windowed=True,
        causal_chunk_length=8,
        causal_window_length=3,
        causal_window_decay=causal_window_decay,
        causal_padding=causal_padding)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    masks = tf.cast(encoder_inputs_mask, dtype=tf.float32)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def test_attention_no_projection(
      self, feature_transform, num_random_features, training, redraw, is_short,
      begin_kernel):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=redraw,
        is_short_seq=is_short,
        begin_kernel=begin_kernel)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    masks = tf.to_float(encoder_inputs_mask)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def test_attention_no_projection(
      self, feature_transform, num_random_features, training, redraw, is_short,
      begin_kernel):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=redraw,
        is_short_seq=is_short,
        begin_kernel=begin_kernel)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    masks = tf.cast(encoder_inputs_mask, dtype=tf.float32)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def test_attention_scale_by_length(self, seq_length):
    num_heads = 12
    key_dim = 64
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        num_random_features=0,
        scale_by_length=True)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.ones((batch_size, seq_length), dtype=tf.int32)
    masks = tf.to_float(encoder_inputs_mask)
    output_scale_by_length = test_layer(
        query=query, value=value, attention_mask=masks)

    test_layer._scale_by_length = False
    output_no_scale_by_length = test_layer(
        query=query, value=value, attention_mask=masks)
    if seq_length == 512:  # Equals because log(seq_length, base=512) = 1.0
      self.assertAllClose(output_scale_by_length, output_no_scale_by_length)
    else:
      self.assertNotAllClose(output_scale_by_length, output_no_scale_by_length)
//  [python]
def test_attention_scale_by_length(self, seq_length):
    num_heads = 12
    key_dim = 64
    batch_size = 2
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        num_random_features=0,
        scale_by_length=True)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    value = query
    encoder_inputs_mask = tf.ones((batch_size, seq_length), dtype=tf.int32)
    masks = tf.cast(encoder_inputs_mask, dtype=tf.float32)
    output_scale_by_length = test_layer(
        query=query, value=value, attention_mask=masks)

    test_layer._scale_by_length = False
    output_no_scale_by_length = test_layer(
        query=query, value=value, attention_mask=masks)
    if seq_length == 512:  # Equals because log(seq_length, base=512) = 1.0
      self.assertAllClose(output_scale_by_length, output_no_scale_by_length)
    else:
      self.assertNotAllClose(output_scale_by_length, output_no_scale_by_length)
//  [python]
def build_mock_model(path):
    """Build and save the mock model with the given signature."""
    module = tf.Module()
    # We have to set this useless variable in order for the TF C API to correctly
    # intake it
    module.var = tf.Variable(0, dtype=tf.float32)

    def action(*inputs):
        s1 = tf.reduce_sum(
            [
                tf.to_float(inputs[0][key])
                for key in PER_LIVEINTERVAL_FEATURE_LIST
            ],
            axis=0,
        )
        s2 = tf.reduce_sum(
            [tf.to_float(inputs[0][key]) for key in CONTEXT_FEATURE_LIST]
        )
        # Add a large number so s won't be 0.
        s = s1 + s2
        result = s + module.var
        return {POLICY_DECISION_LABEL: result}

    module.action = tf.function()(action)
    action = {"action": module.action.get_concrete_function(get_input_signature())}

    tf.saved_model.save(module, path, signatures=action)
    output_spec_path = get_output_spec_path(path)
    with open(output_spec_path, "w") as f:
        print(f"Writing output spec to {output_spec_path}.")
        f.write(POLICY_OUTPUT_SPEC)
//  [python]
def build_mock_model(path):
    """Build and save the mock model with the given signature."""
    module = tf.Module()
    # We have to set this useless variable in order for the TF C API to correctly
    # intake it
    module.var = tf.Variable(0, dtype=tf.float32)

    def action(*inputs):
        s1 = tf.reduce_sum(
            [
                tf.cast(inputs[0][key], tf.float32)
                for key in PER_LIVEINTERVAL_FEATURE_LIST
            ],
            axis=0,
        )
        s2 = tf.reduce_sum(
            [tf.cast(inputs[0][key], tf.float32) for key in CONTEXT_FEATURE_LIST]
        )
        # Add a large number so s won't be 0.
        s = s1 + s2
        result = s + module.var
        return {POLICY_DECISION_LABEL: result}

    module.action = tf.function()(action)
    action = {"action": module.action.get_concrete_function(get_input_signature())}

    tf.saved_model.save(module, path, signatures=action)
    output_spec_path = get_output_spec_path(path)
    with open(output_spec_path, "w") as f:
        print(f"Writing output spec to {output_spec_path}.")
        f.write(POLICY_OUTPUT_SPEC)
//  [python]
def test_causal_windowed_attention_projection_streaming(
      self, feature_transform, causal_chunk_length, causal_weight_decay):
    num_heads = 12
    key_dim = 64
    seq_length = 16
    num_chunks = seq_length // causal_chunk_length
    causal_window_length = num_chunks
    batch_size = 2
    training = False
    num_random_features = 0
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=False,
        is_short_seq=False,
        begin_kernel=False,
        use_causal_windowed=True,
        causal_chunk_length=causal_chunk_length,
        causal_window_length=causal_window_length,
        causal_window_decay=causal_weight_decay,
        causal_padding=None,
        )
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim), seed=2)
    value = query
    encoder_inputs_mask = tf.ones((batch_size, seq_length), dtype=tf.int32)
    masks = tf.to_float(encoder_inputs_mask)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    dim = num_random_features if num_random_features > 0 else key_dim
    kv_cache = tf.zeros(
        (batch_size, num_heads, dim, dim))
    k_sum_cache = tf.zeros((batch_size, num_heads, dim))
    stream_output = []
    cache = {"kv": kv_cache, "k_sum": k_sum_cache}
    for i in range(num_chunks):
      stream_output.append(
          test_layer(
              query=query[:, i * causal_chunk_length:(i + 1) *
                          causal_chunk_length, :],
              value=value[:, i * causal_chunk_length:(i + 1) *
                          causal_chunk_length, :],
              attention_mask=masks[:, i * causal_chunk_length:(i + 1) *
                                   causal_chunk_length],
              cache=cache,
              training=training))
    stream_output = tf.concat(stream_output, axis=1)
    self.assertAllClose(output, stream_output)
//  [python]
def test_causal_windowed_attention_projection_streaming(
      self, feature_transform, causal_chunk_length, causal_weight_decay):
    num_heads = 12
    key_dim = 64
    seq_length = 16
    num_chunks = seq_length // causal_chunk_length
    causal_window_length = num_chunks
    batch_size = 2
    training = False
    num_random_features = 0
    test_layer = attention.KernelAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        feature_transform=feature_transform,
        num_random_features=num_random_features,
        redraw=False,
        is_short_seq=False,
        begin_kernel=False,
        use_causal_windowed=True,
        causal_chunk_length=causal_chunk_length,
        causal_window_length=causal_window_length,
        causal_window_decay=causal_weight_decay,
        causal_padding=None,
        )
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim), seed=2)
    value = query
    encoder_inputs_mask = tf.ones((batch_size, seq_length), dtype=tf.int32)
    masks = tf.cast(encoder_inputs_mask, dtype=tf.float32)
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks,
        training=training)
    dim = num_random_features if num_random_features > 0 else key_dim
    kv_cache = tf.zeros(
        (batch_size, num_heads, dim, dim))
    k_sum_cache = tf.zeros((batch_size, num_heads, dim))
    stream_output = []
    cache = {"kv": kv_cache, "k_sum": k_sum_cache}
    for i in range(num_chunks):
      stream_output.append(
          test_layer(
              query=query[:, i * causal_chunk_length:(i + 1) *
                          causal_chunk_length, :],
              value=value[:, i * causal_chunk_length:(i + 1) *
                          causal_chunk_length, :],
              attention_mask=masks[:, i * causal_chunk_length:(i + 1) *
                                   causal_chunk_length],
              cache=cache,
              training=training))
    stream_output = tf.concat(stream_output, axis=1)
    self.assertAllClose(output, stream_output)
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)

    # 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, ...]

    return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)

    # 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, ...]

    return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def call(self, x, enc_output, training, look_ahead_mask, padding_mask):

        seq_len = tf.shape(x)[1]
        attention_weights = {}

        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
        x *= tf.sqrt(tf.to_float(self.d_model))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x, block1, block2 = self.dec_layers[i](x, enc_output, training, look_ahead_mask, padding_mask)

            attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
            attention_weights['decoder_layer{}_block2'.format(i+1)] = block2

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights
//  [python]
def call(self, x, enc_output, training, look_ahead_mask, padding_mask):

        seq_len = tf.shape(x)[1]
        attention_weights = {}

        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x, block1, block2 = self.dec_layers[i](x, enc_output, training, look_ahead_mask, padding_mask)

            attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
            attention_weights['decoder_layer{}_block2'.format(i+1)] = block2

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights
//  [python]
def indices_to_dense_vector(indices,
                            size,
                            indices_value=1.,
                            default_value=0,
                            dtype=tf.float32):
  """Creates dense vector with indices set to specific value and rest to zeros.

  This function exists because it is unclear if it is safe to use
    tf.sparse_to_dense(indices, [size], 1, validate_indices=False)
  with indices which are not ordered.
  This function accepts a dynamic size (e.g. tf.shape(tensor)[0])

  Args:
    indices: 1d Tensor with integer indices which are to be set to
        indices_values.
    size: scalar with size (integer) of output Tensor.
    indices_value: values of elements specified by indices in the output vector
    default_value: values of other elements in the output vector.
    dtype: data type.

  Returns:
    dense 1D Tensor of shape [size] with indices set to indices_values and the
        rest set to default_value.
  """
  size = t.cast(size, dtype=t.int32)
  zeros = tf.ones([size], dtype=dtype) * default_value
  values = tf.ones_like(indices, dtype=dtype) * indices_value

  return tf.dynamic_stitch(
      [tf.range(size), t.cast(indices, dtype=t.int32)], [zeros, values])
//  [python]
def indices_to_dense_vector(indices,
                            size,
                            indices_value=1.,
                            default_value=0,
                            dtype=tf.float32):
  """Creates dense vector with indices set to specific value and rest to zeros.

  This function exists because it is unclear if it is safe to use
    tf.sparse_to_dense(indices, [size], 1, validate_indices=False)
  with indices which are not ordered.
  This function accepts a dynamic size (e.g. tf.shape(tensor)[0])

  Args:
    indices: 1d Tensor with integer indices which are to be set to
        indices_values.
    size: scalar with size (integer) of output Tensor.
    indices_value: values of elements specified by indices in the output vector
    default_value: values of other elements in the output vector.
    dtype: data type.

  Returns:
    dense 1D Tensor of shape [size] with indices set to indices_values and the
        rest set to default_value.
  """
  size = tf.cast(size, dtype=tf.int32)
  zeros = tf.ones([size], dtype=dtype) * default_value
  values = tf.ones_like(indices, dtype=dtype) * indices_value

  return tf.dynamic_stitch(
      [tf.range(size), tf.cast(indices, dtype=tf.int32)], [zeros, values])
//  [python]
def _log_summaries(input_image, label, num_of_classes, output):
  """Logs the summaries for the model.

  Args:
    input_image: Input image of the model. Its shape is [batch_size, height,
      width, channel].
    label: Label of the image. Its shape is [batch_size, height, width].
    num_of_classes: The number of classes of the dataset.
    output: Output of the model. Its shape is [batch_size, height, width].
  """
  # Add summaries for model variables.
  for model_var in tf.model_variables():
    tf.summary.histogram(model_var.op.name, model_var)

  # Add summaries for images, labels, semantic predictions.
  if FLAGS.save_summaries_images:
    tf.summary.image('samples/%s' % common.IMAGE, input_image)

    # Scale up summary image pixel values for better visualization.
    pixel_scaling = max(1, 255 // num_of_classes)
    summary_label = tf.to_int64(label * pixel_scaling)
    tf.summary.image('samples/%s' % common.LABEL, summary_label)

    predictions = tf.expand_dims(tf.argmax(output, 3), -1)
    summary_predictions = tf.to_int64(predictions * pixel_scaling)
    tf.summary.image('samples/%s' % common.OUTPUT_TYPE, summary_predictions)
//  [python]
def _log_summaries(input_image, label, num_of_classes, output):
  """Logs the summaries for the model.

  Args:
    input_image: Input image of the model. Its shape is [batch_size, height,
      width, channel].
    label: Label of the image. Its shape is [batch_size, height, width].
    num_of_classes: The number of classes of the dataset.
    output: Output of the model. Its shape is [batch_size, height, width].
  """
  # Add summaries for model variables.
  for model_var in tf.model_variables():
    tf.summary.histogram(model_var.op.name, model_var)

  # Add summaries for images, labels, semantic predictions.
  if FLAGS.save_summaries_images:
    tf.summary.image('samples/%s' % common.IMAGE, input_image)

    # Scale up summary image pixel values for better visualization.
    pixel_scaling = max(1, 255 // num_of_classes)
    summary_label = tf.cast(label * pixel_scaling, tf.uint8)
    tf.summary.image('samples/%s' % common.LABEL, summary_label)

    predictions = tf.expand_dims(tf.argmax(output, 3), -1)
    summary_predictions = tf.cast(predictions * pixel_scaling, tf.uint8)
    tf.summary.image('samples/%s' % common.OUTPUT_TYPE, summary_predictions)
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v 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.

    Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable 
          to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
    output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

    # 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)  

    # softmax is 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)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v 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.

    Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable 
          to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
    output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

    # 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)  

    # softmax is 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)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights
//  [python]
def _get_num_pos_neg_samples(self, sorted_indices_tensor, sample_size):
    """Counts the number of positives and negatives numbers to be sampled.

    Args:
      sorted_indices_tensor: A sorted int32 tensor of shape [N] which contains
        the signed indices of the examples where the sign is based on the label
        value. The examples that cannot be sampled are set to 0. It samples
        at most sample_size*positive_fraction positive examples and remaining
        from negative examples.
      sample_size: Size of subsamples.

    Returns:
      A tuple containing the number of positive and negative labels in the
      subsample.
    """
    input_length = tf.shape(input=sorted_indices_tensor)[0]
    valid_positive_index = tf.greater(sorted_indices_tensor,
                                      tf.zeros(input_length, tf.int32))
    num_sampled_pos = tf.reduce_sum(
        input_tensor=tf.to_int32(valid_positive_index))
    max_num_positive_samples = tf.constant(
        int(sample_size * self._positive_fraction), tf.int32)
    num_positive_samples = tf.minimum(max_num_positive_samples, num_sampled_pos)
    num_negative_samples = tf.constant(sample_size,
                                       tf.int32) - num_positive_samples

    return num_positive_samples, num_negative_samples
//  [python]
def _get_num_pos_neg_samples(self, sorted_indices_tensor, sample_size):
    """Counts the number of positives and negatives numbers to be sampled.

    Args:
      sorted_indices_tensor: A sorted int32 tensor of shape [N] which contains
        the signed indices of the examples where the sign is based on the label
        value. The examples that cannot be sampled are set to 0. It samples
        at most sample_size*positive_fraction positive examples and remaining
        from negative examples.
      sample_size: Size of subsamples.

    Returns:
      A tuple containing the number of positive and negative labels in the
      subsample.
    """
    input_length = tf.shape(input=sorted_indices_tensor)[0]
    valid_positive_index = tf.greater(sorted_indices_tensor,
                                      tf.zeros(input_length, tf.int32))
    num_sampled_pos = tf.reduce_sum(
        input_tensor=tf.cast(valid_positive_index, tf.int32))
    max_num_positive_samples = tf.constant(
        int(sample_size * self._positive_fraction), tf.int32)
    num_positive_samples = tf.minimum(max_num_positive_samples, num_sampled_pos)
    num_negative_samples = tf.constant(sample_size,
                                       tf.int32) - num_positive_samples

    return num_positive_samples, num_negative_samples
//  [python]
def test_normalize_image(self, prenormalize, dtype):
    image = tf.constant([[[0, 200, 255]]], dtype=tf.uint8)
    image = tf.tile(image, [64, 64, 1])

    if dtype != tf.uint8 and prenormalize:
      image = image / 255
    image = t.cast(image, dtype=dtype)

    if dtype == tf.uint8 or prenormalize:
      normalized_image = preprocess_ops.normalize_image(
          image, offset=[0.5, 0.5, 0.5], scale=[0.5, 0.5, 0.5]
      )
    else:
      normalized_image = preprocess_ops.normalize_image(
          image, offset=[127.0, 127.0, 127.0], scale=[127.0, 127.0, 127.0]
      )
    max_val = tf.reduce_max(normalized_image)
    # If we mistakely use scale=[0.5, 0.5, 0.5] for non-normalized float input,
    # the normalized image data will contain very large values (e.g. 500).
    tf.assert_greater(2.0, max_val)
//  [python]
def test_normalize_image(self, prenormalize, dtype):
    image = tf.constant([[[0, 200, 255]]], dtype=tf.uint8)
    image = tf.tile(image, [64, 64, 1])

    if dtype != tf.uint8 and prenormalize:
      image = image / 255
    image = tf.cast(image, dtype=dtype)

    if dtype == tf.uint8 or prenormalize:
      normalized_image = preprocess_ops.normalize_image(
          image, offset=[0.5, 0.5, 0.5], scale=[0.5, 0.5, 0.5]
      )
    else:
      normalized_image = preprocess_ops.normalize_image(
          image, offset=[127.0, 127.0, 127.0], scale=[127.0, 127.0, 127.0]
      )
    max_val = tf.reduce_max(normalized_image)
    # If we mistakely use scale=[0.5, 0.5, 0.5] for non-normalized float input,
    # the normalized image data will contain very large values (e.g. 500).
    tf.assert_greater(2.0, max_val)
//  [python]
def _get_values_from_start_and_end(self, input_tensor, num_start_samples,
                                     num_end_samples, total_num_samples):
    """slices num_start_samples and last num_end_samples from input_tensor.

    Args:
      input_tensor: An int32 tensor of shape [N] to be sliced.
      num_start_samples: Number of examples to be sliced from the beginning
        of the input tensor.
      num_end_samples: Number of examples to be sliced from the end of the
        input tensor.
      total_num_samples: Sum of is num_start_samples and num_end_samples. This
        should be a scalar.

    Returns:
      A tensor containing the first num_start_samples and last num_end_samples
      from input_tensor.

    """
    input_length = tf.shape(input=input_tensor)[0]
    start_positions = tf.less(tf.range(input_length), num_start_samples)
    end_positions = tf.greater_equal(
        tf.range(input_length), input_length - num_end_samples)
    selected_positions = tf.logical_or(start_positions, end_positions)
    selected_positions = tf.to_float(selected_positions)
    indexed_positions = tf.multiply(tf.cumsum(selected_positions),
                                    selected_positions)
    one_hot_selector = tf.one_hot(tf.to_int32(indexed_positions) - 1,
                                  total_num_samples,
                                  dtype=tf.float32)
    return tf.to_int32(tf.tensordot(tf.to_float(input_tensor),
                                one_hot_selector, axes=[0, 0]))
//  [python]
def _get_values_from_start_and_end(self, input_tensor, num_start_samples,
                                     num_end_samples, total_num_samples):
    """slices num_start_samples and last num_end_samples from input_tensor.

    Args:
      input_tensor: An int32 tensor of shape [N] to be sliced.
      num_start_samples: Number of examples to be sliced from the beginning
        of the input tensor.
      num_end_samples: Number of examples to be sliced from the end of the
        input tensor.
      total_num_samples: Sum of is num_start_samples and num_end_samples. This
        should be a scalar.

    Returns:
      A tensor containing the first num_start_samples and last num_end_samples
      from input_tensor.

    """
    input_length = tf.shape(input=input_tensor)[0]
    start_positions = tf.less(tf.range(input_length), num_start_samples)
    end_positions = tf.greater_equal(
        tf.range(input_length), input_length - num_end_samples)
    selected_positions = tf.logical_or(start_positions, end_positions)
    selected_positions = tf.cast(selected_positions, tf.float32)
    indexed_positions = tf.multiply(tf.cumsum(selected_positions),
                                    selected_positions)
    one_hot_selector = tf.one_hot(tf.cast(indexed_positions, tf.int32) - 1,
                                  total_num_samples,
                                  dtype=tf.float32)
    return tf.cast(tf.tensordot(tf.cast(input_tensor, tf.float32),
                                one_hot_selector, axes=[0, 0]), tf.int32)
//  [python]
def _get_random_scaled_resolution(
    orig_height,
    orig_width,
    min_scale,
    max_scale,
    max_strech,
    probability_strech):
  """Computes a new random resolution."""
  # Choose a random scale factor and compute new resolution.
  scale = 2 ** tf.random.uniform([],
                                 minval=min_scale,
                                 maxval=max_scale,
                                 dtype=tf.float32)
  scale_height = scale
  scale_width = scale

  # Possibly change scale values individually to perform strech
  def true_fn(scale_height, scale_width):
    scale_height *= 2 ** tf.random.uniform([], -max_strech, max_strech)
    scale_width *= 2 ** tf.random.uniform([], -max_strech, max_strech)
    return tf.stack((scale_height, scale_width), axis=0)
  def false_fn(scale_height, scale_width):
    return tf.stack((scale_height, scale_width), axis=0)
  perform_strech = tf.random.uniform([]) < probability_strech
  scales = tf.cond(perform_strech,
                   lambda: true_fn(scale_height, scale_width),
                   lambda: false_fn(scale_height, scale_width))
  scale_height = scales[0]
  scale_width = scales[1]

  # Compute scaled image resolution.
  new_height = tf.to_int32(
      tf.math.ceil(tf.to_float(orig_height) * scale_height))
  new_width = tf.to_int32(
      tf.math.ceil(tf.to_float(orig_width) * scale_width))
  return new_height, new_width, scale
//  [python]
def _get_random_scaled_resolution(
    orig_height,
    orig_width,
    min_scale,
    max_scale,
    max_strech,
    probability_strech):
  """Computes a new random resolution."""
  # Choose a random scale factor and compute new resolution.
  scale = 2 ** tf.random.uniform([],
                                 minval=min_scale,
                                 maxval=max_scale,
                                 dtype=tf.float32)
  scale_height = scale
  scale_width = scale

  # Possibly change scale values individually to perform strech
  def true_fn(scale_height, scale_width):
    scale_height *= 2 ** tf.random.uniform([], -max_strech, max_strech)
    scale_width *= 2 ** tf.random.uniform([], -max_strech, max_strech)
    return tf.stack((scale_height, scale_width), axis=0)
  def false_fn(scale_height, scale_width):
    return tf.stack((scale_height, scale_width), axis=0)
  perform_strech = tf.random.uniform([]) < probability_strech
  scales = tf.cond(perform_strech,
                   lambda: true_fn(scale_height, scale_width),
                   lambda: false_fn(scale_height, scale_width))
  scale_height = scales[0]
  scale_width = scales[1]

  # Compute scaled image resolution.
  new_height = tf.cast(
      tf.math.ceil(tf.cast(orig_height, tf.float32) * scale_height), tf.int32)
  new_width = tf.cast(
      tf.math.ceil(tf.cast(orig_width, tf.float32) * scale_width), tf.int32)
  return new_height, new_width, scale
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        labels: tf.Tensor | None = None,
        training: bool = False,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states

        outputs = self.wav2vec2(
            input_values,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        if self.config.use_weighted_layer_sum:
            hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
            hidden_states = tf.stack(hidden_states, axis=1)
            norm_weights = tf.nn.softmax(self.layer_weights, axis=-1)
            hidden_states = tf.reduce_sum(hidden_states * tf.reshape(norm_weights, [-1, 1, 1]), axis=1)
        else:
            hidden_states = outputs[0]

        hidden_states = self.projector(hidden_states)
        if attention_mask is None:
            pooled_output = tf.reduce_mean(hidden_states, axis=1)
        else:
            padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask)
            padding_mask_float = tf.cast(padding_mask, hidden_states.dtype)
            hidden_states = tf.multiply(hidden_states, tf.expand_dims(padding_mask_float, axis=-1))
            pooled_output = tf.divide(
                tf.reduce_sum(hidden_states, axis=1), tf.expand_dims(tf.reduce_sum(padding_mask_float, axis=1), axis=1)
            )
        logits = self.classifier(pooled_output)
        loss = None
        if labels is not None:
            loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
            loss = loss_fn(tf.reshape(labels, [-1]), tf.reshape(logits, [-1, self.config.num_labels]))
        if not return_dict:
            output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
            return ((loss,) + output) if loss is not None else output

        return TFSequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
//  [python]
def random_scale_second(images,
                        flow,
                        mask,
                        min_scale,
                        max_scale,
                        probability_scale):
  """Performs a random scaling on the second image in the given range."""
  perform_scale = tf.random.uniform([]) < probability_scale

  def true_fn(images, flow, mask):
    # choose a random scale factor and compute new resolution
    orig_height = tf.shape(images)[-3]
    orig_width = tf.shape(images)[-2]
    new_height, new_width, scale = _get_random_scaled_resolution(
        orig_height=orig_height,
        orig_width=orig_width,
        min_scale=min_scale,
        max_scale=max_scale,
        max_strech=0.0,
        probability_strech=0.0)

    # rescale only the second image
    image_1, image_2 = tf.unstack(images)
    image_2 = smurf_utils.resize(image_2, new_height, new_width, is_flow=False)
    # Crop either first or second image to have matching dimensions
    if scale < 1.0:
      image_1 = _center_crop(image_1, new_height, new_width)
    else:
      image_2 = _center_crop(image_2, orig_height, orig_width)
    images = tf.stack([image_1, image_2])

    if flow is not None:
      # get current locations (with the origin in the image center)
      positions = _positions_center_origin(orig_height, orig_width)

      # compute scale factor of the actual new image resolution
      scale_flow_h = tf.to_float(new_height) / tf.to_float(
          orig_height)
      scale_flow_w = tf.to_float(new_width) / tf.to_float(
          orig_width)
      scale_flow = tf.stack([scale_flow_h, scale_flow_w])

      # compute augmented flow (multiply by mask to zero invalid flow locations)
      flow = ((positions + flow) * scale_flow - positions) * mask

      if scale < 1.0:
        # in case we downsample the image we crop the reference image to keep
        # the same shape
        flow = _center_crop(flow, new_height, new_width)
        mask = _center_crop(mask, new_height, new_width)
    return images, flow, mask
  def false_fn(images, flow, mask):
    return images, flow, mask

  return tf.cond(perform_scale, lambda: true_fn(images, flow, mask),
                 lambda: false_fn(images, flow, mask))
//  [python]
def random_scale_second(images,
                        flow,
                        mask,
                        min_scale,
                        max_scale,
                        probability_scale):
  """Performs a random scaling on the second image in the given range."""
  perform_scale = tf.random.uniform([]) < probability_scale

  def true_fn(images, flow, mask):
    # choose a random scale factor and compute new resolution
    orig_height = tf.shape(images)[-3]
    orig_width = tf.shape(images)[-2]
    new_height, new_width, scale = _get_random_scaled_resolution(
        orig_height=orig_height,
        orig_width=orig_width,
        min_scale=min_scale,
        max_scale=max_scale,
        max_strech=0.0,
        probability_strech=0.0)

    # rescale only the second image
    image_1, image_2 = tf.unstack(images)
    image_2 = smurf_utils.resize(image_2, new_height, new_width, is_flow=False)
    # Crop either first or second image to have matching dimensions
    if scale < 1.0:
      image_1 = _center_crop(image_1, new_height, new_width)
    else:
      image_2 = _center_crop(image_2, orig_height, orig_width)
    images = tf.stack([image_1, image_2])

    if flow is not None:
      # get current locations (with the origin in the image center)
      positions = _positions_center_origin(orig_height, orig_width)

      # compute scale factor of the actual new image resolution
      scale_flow_h = tf.cast(new_height, tf.float32) / tf.cast(
          orig_height, tf.float32)
      scale_flow_w = tf.cast(new_width, tf.float32) / tf.cast(
          orig_width, tf.float32)
      scale_flow = tf.stack([scale_flow_h, scale_flow_w])

      # compute augmented flow (multiply by mask to zero invalid flow locations)
      flow = ((positions + flow) * scale_flow - positions) * mask

      if scale < 1.0:
        # in case we downsample the image we crop the reference image to keep
        # the same shape
        flow = _center_crop(flow, new_height, new_width)
        mask = _center_crop(mask, new_height, new_width)
    return images, flow, mask
  def false_fn(images, flow, mask):
    return images, flow, mask

  return tf.cond(perform_scale, lambda: true_fn(images, flow, mask),
                 lambda: false_fn(images, flow, mask))
//  [python]
def __init__(
            self,
            inputs,
            vocabulary_size,
            embedding_size,
            pad_value=0,
            name='average_embedding_layer',
            embeddings_initializer=tf.random_uniform_initializer(-0.1, 0.1),
            embeddings_kwargs=None,
    ):
        super().__init__(name=name)

        # if embeddings_kwargs is None:
        #     embeddings_kwargs = {}

        if inputs.get_shape().ndims != 2:
            raise ValueError('inputs must be of size batch_size * batch_sentence_length')

        self.inputs = inputs

        print("  [TL] AverageEmbeddingInputlayer %s: (%d, %d)" % (name, vocabulary_size, embedding_size))
        with tf.variable_scope(name):
            self.embeddings = tf.get_variable(
                name='embeddings',
                shape=(vocabulary_size, embedding_size),
                initializer=embeddings_initializer,
                dtype=D_TYPE,
                **(embeddings_kwargs or {})
                # **embeddings_kwargs
            )  # **(embeddings_kwargs or {}),

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

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

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

        self.outputs = sentence_embeddings
        self.all_layers = [self.outputs]
        self.all_params = [self.embeddings]
        self.all_drop = {}
//  [python]
def retrieve_seq_length_op3(data, pad_val=0):  # HangSheng: return tensor for sequence length, if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.reduce_sum(tf.to_int32(tf.reduce_any(tf.not_equal(data, pad_val), axis=2)), 1)
    elif data_shape_size == 2:
        return tf.reduce_sum(tf.to_int32(tf.not_equal(data, pad_val)), 1)
    elif data_shape_size == 1:
        raise ValueError("retrieve_seq_length_op3: data has wrong shape!")
    else:
        raise ValueError("retrieve_seq_length_op3: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))
//  [python]
def retrieve_seq_length_op3(data, pad_val=0):  # HangSheng: return tensor for sequence length, if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.reduce_sum(tf.cast(tf.reduce_any(tf.not_equal(data, pad_val), axis=2), dtype=tf.int32), 1)
    elif data_shape_size == 2:
        return tf.reduce_sum(tf.cast(tf.not_equal(data, pad_val), dtype=tf.int32), 1)
    elif data_shape_size == 1:
        raise ValueError("retrieve_seq_length_op3: data has wrong shape!")
    else:
        raise ValueError("retrieve_seq_length_op3: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))