jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def keypoint_scale(keypoints, y_scale, x_scale, scope=None):
  """Scales keypoint coordinates in x and y dimensions.

  Args:
    keypoints: a tensor of shape [num_instances, num_keypoints, 2]
    y_scale: (float) scalar tensor
    x_scale: (float) scalar tensor
    scope: name scope.

  Returns:
    new_keypoints: a tensor of shape [num_instances, num_keypoints, 2]
  """
  if not scope:
    scope = 'Scale'
  with tf.name_scope(scope):
    y_scale = tf.cast(y_scale, tf.float32)
    x_scale = tf.cast(x_scale, tf.float32)
    new_keypoints = keypoints * [[[y_scale, x_scale]]]
    return new_keypoints
//  [python]
def moller_trumbore(o, d, p0, p1, p2, epsilon):
    r"""
    Computes the intersection between a ray ``ray`` and a triangle defined
    by its vertices ``p0``, ``p1``, and ``p2`` using the Moller–Trumbore
    intersection algorithm.

    Input
    -----
    o, d: [..., 3], tf.float
        Ray origin and direction.
        The direction `d` must be a unit vector.

    p0, p1, p2: [..., 3], tf.float
        Vertices defining the triangle

    epsilon : (), tf.float
        Small value used to avoid errors due to numerical precision

    Output
    -------
    t : [...], tf.float
        Position along the ray from the origin at which the intersection
        occurs (if any)

    hit : [...], bool
        `True` if the ray intersects the triangle. `False` otherwise.
    """

    rdtype = o.dtype
    zero = tf.to_float(0.0, rdtype)
    one = tf.ones((), rdtype)

    # [..., 3]
    e1 = p1 - p0
    e2 = p2 - p0

    # [...,3]
    pvec = cross(d, e2)
    # [...,1]
    det = dot(e1, pvec, keepdim=True)

    # If the ray is parallel to the triangle, then det = 0.
    hit = tf.greater(tf.abs(det), zero)

    # [...,3]
    tvec = o - p0
    # [...,1]
    u = tf.math.divide_no_nan(dot(tvec, pvec, keepdim=True), det)
    # [...,1]
    hit = tf.logical_and(hit,
        tf.logical_and(tf.greater_equal(u, -epsilon),
                       tf.less_equal(u, one + epsilon)))

    # [..., 3]
    qvec = cross(tvec, e1)
    # [...,1]
    v = tf.math.divide_no_nan(dot(d, qvec, keepdim=True), det)
    # [..., 1]
    hit = tf.logical_and(hit,
                            tf.logical_and(tf.greater_equal(v, -epsilon),
                                        tf.less_equal(u + v, one + epsilon)))
    # [..., 1]
    t = tf.math.divide_no_nan(dot(e2, qvec, keepdim=True), det)
    # [..., 1]
    hit = tf.logical_and(hit, tf.greater_equal(t, epsilon))

    # [...]
    t = tf.squeeze(t, axis=-1)
    hit = tf.squeeze(hit, axis=-1)

    return t, hit
//  [python]
def moller_trumbore(o, d, p0, p1, p2, epsilon):
    r"""
    Computes the intersection between a ray ``ray`` and a triangle defined
    by its vertices ``p0``, ``p1``, and ``p2`` using the Moller–Trumbore
    intersection algorithm.

    Input
    -----
    o, d: [..., 3], tf.float
        Ray origin and direction.
        The direction `d` must be a unit vector.

    p0, p1, p2: [..., 3], tf.float
        Vertices defining the triangle

    epsilon : (), tf.float
        Small value used to avoid errors due to numerical precision

    Output
    -------
    t : [...], tf.float
        Position along the ray from the origin at which the intersection
        occurs (if any)

    hit : [...], bool
        `True` if the ray intersects the triangle. `False` otherwise.
    """

    rdtype = o.dtype
    zero = tf.cast(0.0, rdtype)
    one = tf.ones((), rdtype)

    # [..., 3]
    e1 = p1 - p0
    e2 = p2 - p0

    # [...,3]
    pvec = cross(d, e2)
    # [...,1]
    det = dot(e1, pvec, keepdim=True)

    # If the ray is parallel to the triangle, then det = 0.
    hit = tf.greater(tf.abs(det), zero)

    # [...,3]
    tvec = o - p0
    # [...,1]
    u = tf.math.divide_no_nan(dot(tvec, pvec, keepdim=True), det)
    # [...,1]
    hit = tf.logical_and(hit,
        tf.logical_and(tf.greater_equal(u, -epsilon),
                       tf.less_equal(u, one + epsilon)))

    # [..., 3]
    qvec = cross(tvec, e1)
    # [...,1]
    v = tf.math.divide_no_nan(dot(d, qvec, keepdim=True), det)
    # [..., 1]
    hit = tf.logical_and(hit,
                            tf.logical_and(tf.greater_equal(v, -epsilon),
                                        tf.less_equal(u + v, one + epsilon)))
    # [..., 1]
    t = tf.math.divide_no_nan(dot(e2, qvec, keepdim=True), det)
    # [..., 1]
    hit = tf.logical_and(hit, tf.greater_equal(t, epsilon))

    # [...]
    t = tf.squeeze(t, axis=-1)
    hit = tf.squeeze(hit, axis=-1)

    return t, hit
//  [python]
def dataset(directory, images_file, labels_file):
  """Download and parse MNIST dataset."""

  images_file = download(directory, images_file)
  labels_file = download(directory, labels_file)

  check_image_file_header(images_file)
  check_labels_file_header(labels_file)

  def decode_image(image):
    # Normalize from [0, 255] to [0.0, 1.0]
    image = tf.io.decode_raw(image, tf.uint8)
    image = tf.to_float(image)
    image = tf.reshape(image, [784])
    return image / 255.0

  def decode_label(label):
    label = tf.io.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
    label = tf.reshape(label, [])  # label is a scalar
    return tf.to_int32(label)

  images = tf.data.FixedLengthRecordDataset(
      images_file, 28 * 28, header_bytes=16).map(decode_image)
  labels = tf.data.FixedLengthRecordDataset(
      labels_file, 1, header_bytes=8).map(decode_label)
  return tf.data.Dataset.zip((images, labels))
//  [python]
def dataset(directory, images_file, labels_file):
  """Download and parse MNIST dataset."""

  images_file = download(directory, images_file)
  labels_file = download(directory, labels_file)

  check_image_file_header(images_file)
  check_labels_file_header(labels_file)

  def decode_image(image):
    # Normalize from [0, 255] to [0.0, 1.0]
    image = tf.io.decode_raw(image, tf.uint8)
    image = tf.cast(image, tf.float32)
    image = tf.reshape(image, [784])
    return image / 255.0

  def decode_label(label):
    label = tf.io.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
    label = tf.reshape(label, [])  # label is a scalar
    return tf.cast(label, tf.int32)

  images = tf.data.FixedLengthRecordDataset(
      images_file, 28 * 28, header_bytes=16).map(decode_image)
  labels = tf.data.FixedLengthRecordDataset(
      labels_file, 1, header_bytes=8).map(decode_label)
  return tf.data.Dataset.zip((images, labels))
//  [python]
def _compute_new_dynamic_size(image, min_dimension, max_dimension):
  """Compute new dynamic shape for resize_to_range method."""
  image_shape = tf.shape(input=image)
  orig_height = t.cast(image_shape[0], dtype=tf.float32)
  orig_width = t.cast(image_shape[1], dtype=tf.float32)
  num_channels = image_shape[2]
  orig_min_dim = tf.minimum(orig_height, orig_width)
  # Calculates the larger of the possible sizes
  min_dimension = tf.constant(min_dimension, dtype=tf.float32)
  large_scale_factor = min_dimension / orig_min_dim
  # Scaling orig_(height|width) by large_scale_factor will make the smaller
  # dimension equal to min_dimension, save for floating point rounding errors.
  # For reasonably-sized images, taking the nearest integer will reliably
  # eliminate this error.
  large_height = t.cast(
      tf.round(orig_height * large_scale_factor), dtype=tf.int32)
  large_width = t.cast(
      tf.round(orig_width * large_scale_factor), dtype=tf.int32)
  large_size = tf.stack([large_height, large_width])
  if max_dimension:
    # Calculates the smaller of the possible sizes, use that if the larger
    # is too big.
    orig_max_dim = tf.maximum(orig_height, orig_width)
    max_dimension = tf.constant(max_dimension, dtype=tf.float32)
    small_scale_factor = max_dimension / orig_max_dim
    # Scaling orig_(height|width) by small_scale_factor will make the larger
    # dimension equal to max_dimension, save for floating point rounding
    # errors. For reasonably-sized images, taking the nearest integer will
    # reliably eliminate this error.
    small_height = t.cast(
        tf.round(orig_height * small_scale_factor), dtype=tf.int32)
    small_width = t.cast(
        tf.round(orig_width * small_scale_factor), dtype=tf.int32)
    small_size = tf.stack([small_height, small_width])
    new_size = tf.cond(
        pred=t.cast(tf.reduce_max(input_tensor=large_size), dtype=tf.float32) >
        max_dimension,
        true_fn=lambda: small_size,
        false_fn=lambda: large_size)
  else:
    new_size = large_size
  return tf.stack(tf.unstack(new_size) + [num_channels])
//  [python]
def _compute_new_dynamic_size(image, min_dimension, max_dimension):
  """Compute new dynamic shape for resize_to_range method."""
  image_shape = tf.shape(input=image)
  orig_height = tf.cast(image_shape[0], dtype=tf.float32)
  orig_width = tf.cast(image_shape[1], dtype=tf.float32)
  num_channels = image_shape[2]
  orig_min_dim = tf.minimum(orig_height, orig_width)
  # Calculates the larger of the possible sizes
  min_dimension = tf.constant(min_dimension, dtype=tf.float32)
  large_scale_factor = min_dimension / orig_min_dim
  # Scaling orig_(height|width) by large_scale_factor will make the smaller
  # dimension equal to min_dimension, save for floating point rounding errors.
  # For reasonably-sized images, taking the nearest integer will reliably
  # eliminate this error.
  large_height = tf.cast(
      tf.round(orig_height * large_scale_factor), dtype=tf.int32)
  large_width = tf.cast(
      tf.round(orig_width * large_scale_factor), dtype=tf.int32)
  large_size = tf.stack([large_height, large_width])
  if max_dimension:
    # Calculates the smaller of the possible sizes, use that if the larger
    # is too big.
    orig_max_dim = tf.maximum(orig_height, orig_width)
    max_dimension = tf.constant(max_dimension, dtype=tf.float32)
    small_scale_factor = max_dimension / orig_max_dim
    # Scaling orig_(height|width) by small_scale_factor will make the larger
    # dimension equal to max_dimension, save for floating point rounding
    # errors. For reasonably-sized images, taking the nearest integer will
    # reliably eliminate this error.
    small_height = tf.cast(
        tf.round(orig_height * small_scale_factor), dtype=tf.int32)
    small_width = tf.cast(
        tf.round(orig_width * small_scale_factor), dtype=tf.int32)
    small_size = tf.stack([small_height, small_width])
    new_size = tf.cond(
        pred=tf.cast(tf.reduce_max(input_tensor=large_size), dtype=tf.float32) >
        max_dimension,
        true_fn=lambda: small_size,
        false_fn=lambda: large_size)
  else:
    new_size = large_size
  return tf.stack(tf.unstack(new_size) + [num_channels])
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    predicted_values, _ = self._constraint_network(observation, training=False)
    is_satisfied = self._comparator_fn(predicted_values, self._absolute_value)
    return t.cast(is_satisfied, t.float32)
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    predicted_values, _ = self._constraint_network(observation, training=False)
    is_satisfied = self._comparator_fn(predicted_values, self._absolute_value)
    return tf.cast(is_satisfied, tf.float32)
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    predicted_quantiles, _ = self._constraint_network(
        observation, training=False
    )
    is_satisfied = self._comparator_fn(
        predicted_quantiles, self._quantile_value
    )
    return t.cast(is_satisfied, t.float32)
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    predicted_quantiles, _ = self._constraint_network(
        observation, training=False
    )
    is_satisfied = self._comparator_fn(
        predicted_quantiles, self._quantile_value
    )
    return tf.cast(is_satisfied, tf.float32)
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    predicted_values, _ = self._constraint_network(observation, training=False)

    batch_dims = nest_utils.get_outer_shape(
        observation, self._time_step_spec.observation
    )
    if self._baseline_action_fn is not None:
      baseline_action = self._baseline_action_fn(observation)
      baseline_action.shape.assert_is_compatible_with(batch_dims)
    else:
      baseline_action = tf.zeros(batch_dims, dtype=tf.int32)

    predicted_values_for_baseline_actions = common.index_with_actions(
        predicted_values, tf.cast(baseline_action, dtype=tf.int32)
    )
    predicted_values_for_baseline_actions = self._reshape_and_broadcast(
        predicted_values_for_baseline_actions, tf.shape(predicted_values)
    )
    is_satisfied = self._comparator_fn(
        predicted_values,
        (1 - self._margin) * predicted_values_for_baseline_actions,
    )
    return tf.cast(is_satisfied, tf.float32)
//  [python]
def compute_loss(
      self,
      observations: types.NestedTensor,
      actions: types.NestedTensor,
      rewards: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> types.Tensor:
    """Computes loss for training the constraint network.

    Args:
      observations: A batch of observations.
      actions: A batch of actions.
      rewards: A batch of rewards.
      weights: Optional scalar or elementwise (per-batch-entry) importance
        weights.  The output batch loss will be scaled by these weights, and the
        final scalar loss is the mean of these values.
      training: Whether the loss is being used for training.

    Returns:
      loss: A `Tensor` containing the loss for the training step.
    """
    with tf.name_scope('constraint_loss'):
      sample_weights = weights if weights is not None else 1
      predicted_values, _ = self._constraint_network(
          observations, training=training
      )
      action_predicted_values = common.index_with_actions(
          predicted_values, tf.cast(actions, dtype=tf.int32)
      )
      # Reduction is done outside of the loss function because non-scalar
      # weights with unknown shapes may trigger shape validation that fails
      # XLA compilation.
      return tf.reduce_mean(
          tf.multiply(
              self._error_loss_fn(
                  rewards,
                  action_predicted_values,
                  reduction=tf.compat.v1.losses.Reduction.NONE,
              ),
              sample_weights,
          )
      )
//  [python]
def get_masked_bias(dtype: tf.DType) -> tf.Tensor:
        return tf.cast(tf.constant(-1e9), dtype)
//  [python]
def get_causal_mask(self, key_length, query_length) -> tf.Tensor:
        return tf.to_float(self.lower_triangle_mask[:, :, key_length - query_length : key_length, :key_length])
//  [python]
def get_causal_mask(self, key_length, query_length) -> tf.Tensor:
        return tf.cast(self.lower_triangle_mask[:, :, key_length - query_length : key_length, :key_length], tf.bool)
//  [python]
def create_sinusoidal_positions(num_pos: int, dim: int) -> tf.Tensor:
    inv_freq = t.cast(1.0 / (10000 ** (t.range(0, dim, 2) / dim)), t.float32)
    sinusoid_inp = t.cast(t.einsum("i , j -> i j", t.range(num_pos, dtype=t.float32), inv_freq), t.float32)
    sin, cos = t.sin(sinusoid_inp), t.cos(sinusoid_inp)
    out = t.concat((sin, cos), axis=1)
    return out
//  [python]
def create_sinusoidal_positions(num_pos: int, dim: int) -> tf.Tensor:
    inv_freq = tf.cast(1.0 / (10000 ** (tf.range(0, dim, 2) / dim)), tf.float32)
    sinusoid_inp = tf.cast(tf.einsum("i , j -> i j", tf.range(num_pos, dtype=tf.float32), inv_freq), tf.float32)
    sin, cos = tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)
    out = tf.concat((sin, cos), axis=1)
    return out
//  [python]
def compute_feasibility_probability(
    observation: types.NestedTensor,
    constraints: Iterable[BaseConstraint],
    batch_size: types.Int,
    num_actions: int,
    action_mask: Optional[types.Tensor] = None,
) -> types.Float:
  """Helper function to compute the action feasibility probability."""
  feasibility_prob = tf.ones([batch_size, num_actions])
  if action_mask is not None:
    feasibility_prob = t.cast(action_mask, tf.float32)
  for c in constraints:
    # We assume the constraints are independent.
    action_feasibility = c(observation)
    feasibility_prob *= action_feasibility
  return feasibility_prob
//  [python]
def compute_feasibility_probability(
    observation: types.NestedTensor,
    constraints: Iterable[BaseConstraint],
    batch_size: types.Int,
    num_actions: int,
    action_mask: Optional[types.Tensor] = None,
) -> types.Float:
  """Helper function to compute the action feasibility probability."""
  feasibility_prob = tf.ones([batch_size, num_actions])
  if action_mask is not None:
    feasibility_prob = tf.cast(action_mask, tf.float32)
  for c in constraints:
    # We assume the constraints are independent.
    action_feasibility = c(observation)
    feasibility_prob *= action_feasibility
  return feasibility_prob
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    predicted_quantiles, _ = self._constraint_network(
        observation, training=False
    )
    batch_dims = nest_utils.get_outer_shape(
        observation, self._time_step_spec.observation
    )

    if self._baseline_action_fn is not None:
      baseline_action = self._baseline_action_fn(observation)
      baseline_action.shape.assert_is_compatible_with(batch_dims)
    else:
      baseline_action = tf.zeros(batch_dims, dtype=tf.int32)

    predicted_quantiles_for_baseline_actions = common.index_with_actions(
        predicted_quantiles, tf.cast(baseline_action, dtype=tf.int32)
    )
    predicted_quantiles_for_baseline_actions = self._reshape_and_broadcast(
        predicted_quantiles_for_baseline_actions, tf.shape(predicted_quantiles)
    )
    is_satisfied = self._comparator_fn(
        predicted_quantiles, predicted_quantiles_for_baseline_actions
    )
    return t.cast(is_satisfied, tf.float32)
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    predicted_quantiles, _ = self._constraint_network(
        observation, training=False
    )
    batch_dims = nest_utils.get_outer_shape(
        observation, self._time_step_spec.observation
    )

    if self._baseline_action_fn is not None:
      baseline_action = self._baseline_action_fn(observation)
      baseline_action.shape.assert_is_compatible_with(batch_dims)
    else:
      baseline_action = tf.zeros(batch_dims, dtype=tf.int32)

    predicted_quantiles_for_baseline_actions = common.index_with_actions(
        predicted_quantiles, tf.cast(baseline_action, dtype=tf.int32)
    )
    predicted_quantiles_for_baseline_actions = self._reshape_and_broadcast(
        predicted_quantiles_for_baseline_actions, tf.shape(predicted_quantiles)
    )
    is_satisfied = self._comparator_fn(
        predicted_quantiles, predicted_quantiles_for_baseline_actions
    )
    return tf.cast(is_satisfied, tf.float32)
//  [python]
def covariance_loss(self):

        X = tf.concat((self.item_embeddings, self.user_embeddings), 0)
        n_rows = tf.to_float(tf.shape(X)[0])
        X = X - (tf.reduce_mean(X, axis=0))
        cov = tf.matmul(X, X, transpose_a=True) / n_rows

        return tf.reduce_sum(tf.matrix_set_diag(cov, tf.zeros(self.embed_dim, tf.float32))) * self.cov_loss_weight
//  [python]
def covariance_loss(self):

        X = tf.concat((self.item_embeddings, self.user_embeddings), 0)
        n_rows = tf.cast(tf.shape(X)[0], tf.float32)
        X = X - (tf.reduce_mean(X, axis=0))
        cov = tf.matmul(X, X, transpose_a=True) / n_rows

        return tf.reduce_sum(tf.matrix_set_diag(cov, tf.zeros(self.embed_dim, tf.float32))) * self.cov_loss_weight
//  [python]
def _attn(
        self,
        query: tf.Tensor,
        key: tf.Tensor,
        value: tf.Tensor,
        attention_mask: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        # compute causal mask from causal mask buffer
        query_length, key_length = shape_list(query)[-2], shape_list(key)[-2]
        causal_mask = self.get_causal_mask(key_length, query_length)

        # Keep the attention weights computation in fp32 to avoid overflow issues
        query = t.cast(query, tf.float32)
        key = t.cast(key, tf.float32)

        attn_weights = tf.matmul(query, key, transpose_b=True)
        attn_weights = tf.where(causal_mask, attn_weights, self.get_masked_bias(attn_weights.dtype))

        attn_weights = attn_weights / self.scale_attn

        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        attn_weights = stable_softmax(attn_weights, axis=-1)
        attn_weights = t.cast(attn_weights, value.dtype)
        attn_weights = self.attn_dropout(attn_weights)

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

        attn_output = tf.matmul(attn_weights, value)

        return attn_output, attn_weights
//  [python]
def _attn(
        self,
        query: tf.Tensor,
        key: tf.Tensor,
        value: tf.Tensor,
        attention_mask: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        # compute causal mask from causal mask buffer
        query_length, key_length = shape_list(query)[-2], shape_list(key)[-2]
        causal_mask = self.get_causal_mask(key_length, query_length)

        # Keep the attention weights computation in fp32 to avoid overflow issues
        query = tf.cast(query, tf.float32)
        key = tf.cast(key, tf.float32)

        attn_weights = tf.matmul(query, key, transpose_b=True)
        attn_weights = tf.where(causal_mask, attn_weights, self.get_masked_bias(attn_weights.dtype))

        attn_weights = attn_weights / self.scale_attn

        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        attn_weights = stable_softmax(attn_weights, axis=-1)
        attn_weights = tf.cast(attn_weights, value.dtype)
        attn_weights = self.attn_dropout(attn_weights)

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

        attn_output = tf.matmul(attn_weights, value)

        return attn_output, attn_weights
//  [python]
def result(self) -> Dict[str, tf.Tensor]:
    """Returns the metrics values as a dict."""
    # (num_categories,)
    tp_fn_fp_count = self.tp_count + self.fn_count + self.fp_count
    is_ignore_label = tf.one_hot(
        self._ignored_label,
        self._num_categories,
        on_value=True,
        off_value=False,
    )

    sq_per_class = tf.math.divide_no_nan(
        self.tp_iou_sum, self.tp_count
    ) * tf.to_float(~is_ignore_label)
    rq_per_class = tf.math.divide_no_nan(
        self.tp_count, self.tp_count + 0.5 * self.fp_count + 0.5 * self.fn_count
    ) * tf.to_float(~is_ignore_label)
    pq_per_class = sq_per_class * rq_per_class
    is_thing = tf.constant(self._is_thing, dtype=tf.bool)

    result = {
        # (num_categories,)
        'valid_thing_classes': (
            (tp_fn_fp_count > 0) & is_thing & ~is_ignore_label
        ),
        # (num_categories,)
        'valid_stuff_classes': (
            (tp_fn_fp_count > 0) & ~is_thing & ~is_ignore_label
        ),
        # (num_categories,)
        'sq_per_class': sq_per_class,
        # (num_categories,)
        'rq_per_class': rq_per_class,
        # (num_categories,)
        'pq_per_class': pq_per_class,
    }
    return result
//  [python]
def result(self) -> Dict[str, tf.Tensor]:
    """Returns the metrics values as a dict."""
    # (num_categories,)
    tp_fn_fp_count = self.tp_count + self.fn_count + self.fp_count
    is_ignore_label = tf.one_hot(
        self._ignored_label,
        self._num_categories,
        on_value=True,
        off_value=False,
    )

    sq_per_class = tf.math.divide_no_nan(
        self.tp_iou_sum, self.tp_count
    ) * tf.cast(~is_ignore_label, tf.float32)
    rq_per_class = tf.math.divide_no_nan(
        self.tp_count, self.tp_count + 0.5 * self.fp_count + 0.5 * self.fn_count
    ) * tf.cast(~is_ignore_label, tf.float32)
    pq_per_class = sq_per_class * rq_per_class
    is_thing = tf.constant(self._is_thing, dtype=tf.bool)

    result = {
        # (num_categories,)
        'valid_thing_classes': (
            (tp_fn_fp_count > 0) & is_thing & ~is_ignore_label
        ),
        # (num_categories,)
        'valid_stuff_classes': (
            (tp_fn_fp_count > 0) & ~is_thing & ~is_ignore_label
        ),
        # (num_categories,)
        'sq_per_class': sq_per_class,
        # (num_categories,)
        'rq_per_class': rq_per_class,
        # (num_categories,)
        'pq_per_class': pq_per_class,
    }
    return result
//  [python]
def scaled_dot_product_attention(q, k, v, mask, is_training=True, rel_key_emb=None, rel_val_emb=None):
        # 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)

        # attention_weights = tf.layers.dropout(attention_weights, training=is_training, rate=0.2)
        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, is_training=True, rel_key_emb=None, rel_val_emb=None):
        # 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)

        # attention_weights = tf.layers.dropout(attention_weights, training=is_training, rate=0.2)
        output = tf.matmul(attention_weights, v)  # (..., num_heads, attn_dim, depth)

        return output, attention_weights
//  [python]
def create_padding_mask(seq):
  seq = tf.to_float(tf.math.equal(seq, 0))
  
  # add extra dimensions to add the padding
  # to the attention logits.
  return seq[:, tf.newaxis, tf.newaxis, :]
//  [python]
def create_padding_mask(seq):
  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)
  
  # add extra dimensions to add the padding
  # to the attention logits.
  return seq[:, tf.newaxis, tf.newaxis, :]
//  [python]
def update_state(
      self, y_true: Dict[str, tf.Tensor], y_pred: Dict[str, tf.Tensor]
  ):
    category_mask = tf.convert_to_tensor(y_pred['category_mask'], tf.int32)
    instance_mask = tf.convert_to_tensor(y_pred['instance_mask'], tf.int32)
    gt_category_mask = tf.convert_to_tensor(y_true['category_mask'], tf.int32)
    gt_instance_mask = tf.convert_to_tensor(y_true['instance_mask'], tf.int32)

    if self._rescale_predictions:
      _, height, width = gt_category_mask.get_shape().as_list()
      # Instead of cropping the masks to the original image shape (dynamic),
      # here we keep the mask shape (fixed) and ignore the pixels outside the
      # original image shape.
      image_shape = t.cast(y_true['image_info'][:, 0, :], tf.int32)
      # (batch_size, 2)
      y0_x0 = tf.broadcast_to(
          tf.constant([[0, 0]], dtype=tf.int32), tf.shape(image_shape)
      )
      # (batch_size, 4)
      image_shape_bbox = tf.concat([y0_x0, image_shape], axis=1)
      # (batch_size, height, width)
      image_shape_masks = box_ops.bbox2mask(
          bbox=image_shape_bbox,
          image_height=height,
          image_width=width,
          dtype=tf.bool,
      )
      # (batch_size, height, width)
      category_mask = tf.where(
          image_shape_masks, category_mask, self._ignored_label
      )
      instance_mask = tf.where(image_shape_masks, instance_mask, 0)
      gt_category_mask = tf.where(
          image_shape_masks, gt_category_mask, self._ignored_label
      )
      gt_instance_mask = tf.where(image_shape_masks, gt_instance_mask, 0)

    self._update_thing_classes(
        category_mask, instance_mask, gt_category_mask, gt_instance_mask
    )
    self._update_stuff_classes(category_mask, gt_category_mask)
//  [python]
def update_state(
      self, y_true: Dict[str, tf.Tensor], y_pred: Dict[str, tf.Tensor]
  ):
    category_mask = tf.convert_to_tensor(y_pred['category_mask'], tf.int32)
    instance_mask = tf.convert_to_tensor(y_pred['instance_mask'], tf.int32)
    gt_category_mask = tf.convert_to_tensor(y_true['category_mask'], tf.int32)
    gt_instance_mask = tf.convert_to_tensor(y_true['instance_mask'], tf.int32)

    if self._rescale_predictions:
      _, height, width = gt_category_mask.get_shape().as_list()
      # Instead of cropping the masks to the original image shape (dynamic),
      # here we keep the mask shape (fixed) and ignore the pixels outside the
      # original image shape.
      image_shape = tf.cast(y_true['image_info'][:, 0, :], tf.int32)
      # (batch_size, 2)
      y0_x0 = tf.broadcast_to(
          tf.constant([[0, 0]], dtype=tf.int32), tf.shape(image_shape)
      )
      # (batch_size, 4)
      image_shape_bbox = tf.concat([y0_x0, image_shape], axis=1)
      # (batch_size, height, width)
      image_shape_masks = box_ops.bbox2mask(
          bbox=image_shape_bbox,
          image_height=height,
          image_width=width,
          dtype=tf.bool,
      )
      # (batch_size, height, width)
      category_mask = tf.where(
          image_shape_masks, category_mask, self._ignored_label
      )
      instance_mask = tf.where(image_shape_masks, instance_mask, 0)
      gt_category_mask = tf.where(
          image_shape_masks, gt_category_mask, self._ignored_label
      )
      gt_instance_mask = tf.where(image_shape_masks, gt_instance_mask, 0)

    self._update_thing_classes(
        category_mask, instance_mask, gt_category_mask, gt_instance_mask
    )
    self._update_stuff_classes(category_mask, gt_category_mask)
//  [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,
        hidden_states: tf.Tensor,
        layer_past: Optional[Tuple[tf.Tensor, tf.Tensor]] = None,
        attention_mask: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
    ):
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = self._split_heads(query, True)
        key = self._split_heads(key, True)
        value = self._split_heads(value, False)

        sincos = t.cast(t.gather(self.embed_positions, position_ids, axis=0), hidden_states.dtype)
        sincos = tf.split(sincos, 2, axis=-1)
        if self.rotary_dim is not None:
            k_rot = key[:, :, :, : self.rotary_dim]
            k_pass = key[:, :, :, self.rotary_dim :]

            q_rot = query[:, :, :, : self.rotary_dim]
            q_pass = query[:, :, :, self.rotary_dim :]

            k_rot = apply_rotary_pos_emb(k_rot, sincos)
            q_rot = apply_rotary_pos_emb(q_rot, sincos)

            key = tf.concat((k_rot, k_pass), axis=-1)
            query = tf.concat((q_rot, q_pass), axis=-1)
        else:
            key = apply_rotary_pos_emb(key, sincos)
            query = apply_rotary_pos_emb(query, sincos)

        key = tf.transpose(key, (0, 2, 1, 3))
        query = tf.transpose(query, (0, 2, 1, 3))

        if layer_past is not None:
            past_key = layer_past[0]
            past_value = layer_past[1]
            key = tf.concat((past_key, key), axis=-2)
            value = tf.concat((past_value, value), axis=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        # compute self-attention: V x Softmax(QK^T)
        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output)
        attn_output = self.out_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        layer_past: Optional[Tuple[tf.Tensor, tf.Tensor]] = None,
        attention_mask: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
    ):
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = self._split_heads(query, True)
        key = self._split_heads(key, True)
        value = self._split_heads(value, False)

        sincos = tf.cast(tf.gather(self.embed_positions, position_ids, axis=0), hidden_states.dtype)
        sincos = tf.split(sincos, 2, axis=-1)
        if self.rotary_dim is not None:
            k_rot = key[:, :, :, : self.rotary_dim]
            k_pass = key[:, :, :, self.rotary_dim :]

            q_rot = query[:, :, :, : self.rotary_dim]
            q_pass = query[:, :, :, self.rotary_dim :]

            k_rot = apply_rotary_pos_emb(k_rot, sincos)
            q_rot = apply_rotary_pos_emb(q_rot, sincos)

            key = tf.concat((k_rot, k_pass), axis=-1)
            query = tf.concat((q_rot, q_pass), axis=-1)
        else:
            key = apply_rotary_pos_emb(key, sincos)
            query = apply_rotary_pos_emb(query, sincos)

        key = tf.transpose(key, (0, 2, 1, 3))
        query = tf.transpose(query, (0, 2, 1, 3))

        if layer_past is not None:
            past_key = layer_past[0]
            past_value = layer_past[1]
            key = tf.concat((past_key, key), axis=-2)
            value = tf.concat((past_value, value), axis=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        # compute self-attention: V x Softmax(QK^T)
        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output)
        attn_output = self.out_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs
//  [python]
def embeddings_chain(batch_size, num_embs, dims, sigma):
  tri_int = t.cast(num_embs * (num_embs + 1) / 2, dtype=t.int32)
  lin_op = tfp.math.fill_triangular(tf.ones(tri_int))
  lin_op_repeat = tf.tile(lin_op[tf.newaxis, tf.newaxis, Ellipsis],
                          multiples=[batch_size, dims, 1, 1])
  initial_vec = tf.random.uniform((batch_size, 1, dims), minval=-1, maxval=1)
  noise = sigma * tf.random.normal((batch_size, num_embs - 1, dims))
  initial_and_noise = tf.transpose(tf.concat([initial_vec, noise], axis=1),
                                   (0, 2, 1))
  emb_chain = tf.matmul(lin_op_repeat, initial_and_noise[Ellipsis, tf.newaxis])
  return tf.squeeze(tf.transpose(emb_chain, (0, 2, 1, 3)))
//  [python]
def embeddings_chain(batch_size, num_embs, dims, sigma):
  tri_int = tf.cast(num_embs * (num_embs + 1) / 2, dtype=tf.int32)
  lin_op = tfp.math.fill_triangular(tf.ones(tri_int))
  lin_op_repeat = tf.tile(lin_op[tf.newaxis, tf.newaxis, Ellipsis],
                          multiples=[batch_size, dims, 1, 1])
  initial_vec = tf.random.uniform((batch_size, 1, dims), minval=-1, maxval=1)
  noise = sigma * tf.random.normal((batch_size, num_embs - 1, dims))
  initial_and_noise = tf.transpose(tf.concat([initial_vec, noise], axis=1),
                                   (0, 2, 1))
  emb_chain = tf.matmul(lin_op_repeat, initial_and_noise[Ellipsis, tf.newaxis])
  return tf.squeeze(tf.transpose(emb_chain, (0, 2, 1, 3)))
//  [python]
def __init__(self, d_model, warmup_steps=4000):
    super(CustomSchedule, self).__init__()
    
    self.d_model = d_model
    self.d_model = tf.cast(self.d_model, tf.float32)

    self.warmup_steps = warmup_steps
//  [python]
def loss_function(real, pred):
  mask = tf.math.logical_not(tf.math.equal(real, 0))
  loss_ = loss_object(real, pred)

  mask = tf.to_float(mask)
  loss_ *= mask
  
  return tf.reduce_sum(loss_)/tf.reduce_sum(mask)
//  [python]
def loss_function(real, pred):
  mask = tf.math.logical_not(tf.math.equal(real, 0))
  loss_ = loss_object(real, pred)

  mask = tf.cast(mask, dtype=loss_.dtype)
  loss_ *= mask
  
  return tf.reduce_sum(loss_)/tf.reduce_sum(mask)
//  [python]
def __call__(self, x):
        return t.cast(x, dtype=self.dtype)
//  [python]
def __call__(self, x):
        return tf.cast(x, dtype=self.dtype)
//  [python]
def cast(x, dtype):
    """
    Casts a tensor to a new type.

    Parameters
    ----------
    x : tensor
        A Tensor or SparseTensor or IndexedSlices of numeric type.
        It could be uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64.
    dtype : dtpye
         The destination type. The list of supported dtypes is the same as x

    Returns
    -------
        A Tensor or SparseTensor or IndexedSlices with same shape as x and same type as dtype.
    """

    return t.cast(x, dtype=dtype)
//  [python]
def cast(x, dtype):
    """
    Casts a tensor to a new type.

    Parameters
    ----------
    x : tensor
        A Tensor or SparseTensor or IndexedSlices of numeric type.
        It could be uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64.
    dtype : dtpye
         The destination type. The list of supported dtypes is the same as x

    Returns
    -------
        A Tensor or SparseTensor or IndexedSlices with same shape as x and same type as dtype.
    """

    return tf.cast(x, dtype=dtype)
//  [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])   ## 64
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

  # add the mask to the scaled tensor.
  if mask is not None:
    scaled_attention_logits = 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)   ## 64
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

  # add the mask to the scaled tensor.
  if mask is not None:
    scaled_attention_logits = 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 parse_tfrecords_function(example_proto):
  """
    Decode TFRecords for Dataset.
    
    Args:
      example_proto: TensorFlow ExampleProto object. 
    
    Return:
      The op of features and labels
    """

  if FLAGS.label_type == "int":
    features = {
        "ids": tf.VarLenFeature(tf.int64),
        "values": tf.VarLenFeature(tf.float32),
        "label": tf.FixedLenFeature([], tf.int64, default_value=0)
    }

    parsed_features = tf.parse_single_example(example_proto, features)
    labels = parsed_features["label"]
    ids = parsed_features["ids"]
    values = parsed_features["values"]

  elif FLAGS.label_type == "float":
    features = {
        "ids": tf.VarLenFeature(tf.int64),
        "values": tf.VarLenFeature(tf.float32),
        "label": tf.FixedLenFeature([], tf.float32, default_value=0)
    }

    parsed_features = tf.parse_single_example(example_proto, features)
    labels = tf.to_int32(parsed_features["label"])
    ids = parsed_features["ids"]
    values = parsed_features["values"]

  return labels, ids, values
//  [python]
def parse_tfrecords_function(example_proto):
  """
    Decode TFRecords for Dataset.
    
    Args:
      example_proto: TensorFlow ExampleProto object. 
    
    Return:
      The op of features and labels
    """

  if FLAGS.label_type == "int":
    features = {
        "ids": tf.VarLenFeature(tf.int64),
        "values": tf.VarLenFeature(tf.float32),
        "label": tf.FixedLenFeature([], tf.int64, default_value=0)
    }

    parsed_features = tf.parse_single_example(example_proto, features)
    labels = parsed_features["label"]
    ids = parsed_features["ids"]
    values = parsed_features["values"]

  elif FLAGS.label_type == "float":
    features = {
        "ids": tf.VarLenFeature(tf.int64),
        "values": tf.VarLenFeature(tf.float32),
        "label": tf.FixedLenFeature([], tf.float32, default_value=0)
    }

    parsed_features = tf.parse_single_example(example_proto, features)
    labels = tf.cast(parsed_features["label"], tf.int32)
    ids = parsed_features["ids"]
    values = parsed_features["values"]

  return labels, ids, values
//  [python]
def evaluate(inp_sentence):
  

  en_START_TOKEN_id = en_dictionary['<sos>']
  en_END_TOKEN_id   = en_dictionary['<eos>']
    
  pt_START_TOKEN_id = pt_dictionary['<sos>']
  pt_END_TOKEN_id   = pt_dictionary['<eos>']

  start_token = [pt_START_TOKEN_id]
  end_token   = [pt_END_TOKEN_id]

  
  # inp sentence is portuguese, hence adding the start and end token
  inp_sentence = start_token + encode(inp_sentence, pt_dictionary) + end_token
  encoder_input = tf.expand_dims(inp_sentence, 0)
  
  # as the target is english, the first word to the transformer should be the
  # english start token.
  decoder_input = [en_START_TOKEN_id]
  output = tf.expand_dims(decoder_input, 0)
    
  for i in range(MAX_LENGTH):
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(encoder_input, output)
  
    # predictions.shape == (batch_size, seq_len, vocab_size)
    predictions, attention_weights = transformer(encoder_input, 
                                                 output,
                                                 False,
                                                 enc_padding_mask,
                                                 combined_mask,
                                                 dec_padding_mask)
    
    # select the last word from the seq_len dimension
    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

    predicted_id = tf.to_int32(tf.argmax(predictions, axis=-1))
    
    # return the result if the predicted_id is equal to the end token
    if predicted_id == en_END_TOKEN_id:
      return tf.squeeze(output, axis=0), attention_weights
    
    # concatentate the predicted_id to the output which is given to the decoder
    # as its input.
    output = tf.concat([output, predicted_id], axis=-1)

  return tf.squeeze(output, axis=0), attention_weights