jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def self_clustering(
    representations: tf.Tensor, *, subtract_mean: bool = False, **_
) -> tf.Tensor:
  """Self-clustering metric implementation.

  Computes a metric that measures how well distributed representations are, if
  projected on the unit sphere. If `subtract_mean` is True, we additionally
  remove the mean from representations. The metric has a range of
  (-0.5, 1]. It achieves its maximum of 1 if representations collapse to a
  single point, and it is approximately 0 if representations are distributed
  randomly on the sphere. In theory, it can achieve negative values if the
  points are maximally equiangular, although this is very rare in practice.
  Refer to https://arxiv.org/abs/2305.16562 for more details.

  Args:
    representations: Input representations.
    subtract_mean: Whether to subtract the mean from representations.

  Returns:
    Metric value as scalar `tf.Tensor`.
  """
  if representations.shape.rank != 2:
    raise ValueError(f"Expected 2D tensor (got shape {representations.shape})")
  if subtract_mean:
    representations -= tf.reduce_mean(representations, axis=0)
  representations, _ = tf.linalg.normalize(representations, ord=2, axis=1)
  batch_size, feature_dim = tf.unstack(
      tf.cast(tf.shape(representations), tf.float32)
  )
  if tf.math.reduce_all(tf.math.is_nan(representations)):
    # If representations are completely collapsed, return 1.
    return tf.constant(1.0)
  expected = batch_size + batch_size * (batch_size - 1) / feature_dim
  actual = tf.reduce_sum(
      tf.square(tf.matmul(representations, representations, transpose_b=True))
  )
  return (actual - expected) / (batch_size * batch_size - expected)
//  [python]
def draw_side_by_side_evaluation_image(eval_dict,
                                       category_index,
                                       max_boxes_to_draw=20,
                                       min_score_thresh=0.2,
                                       use_normalized_coordinates=True):
  """Creates a side-by-side image with detections and groundtruth.

  Bounding boxes (and instance masks, if available) are visualized on both
  subimages.

  Args:
    eval_dict: The evaluation dictionary returned by
      eval_util.result_dict_for_single_example().
    category_index: A category index (dictionary) produced from a labelmap.
    max_boxes_to_draw: The maximum number of boxes to draw for detections.
    min_score_thresh: The minimum score threshold for showing detections.
    use_normalized_coordinates: Whether to assume boxes and kepoints are in
      normalized coordinates (as opposed to absolute coordiantes).
      Default is True.

  Returns:
    A [1, H, 2 * W, C] uint8 tensor. The subimage on the left corresponds to
      detections, while the subimage on the right corresponds to groundtruth.
  """
  detection_fields = fields.DetectionResultFields()
  input_data_fields = fields.InputDataFields()
  instance_masks = None
  if detection_fields.detection_masks in eval_dict:
    instance_masks = t.cast(
        t.expand_dims(eval_dict[detection_fields.detection_masks], axis=0),
        t.uint8)
  keypoints = None
  if detection_fields.detection_keypoints in eval_dict:
    keypoints = t.expand_dims(
        eval_dict[detection_fields.detection_keypoints], axis=0)
  groundtruth_instance_masks = None
  if input_data_fields.groundtruth_instance_masks in eval_dict:
    groundtruth_instance_masks = t.cast(
        t.expand_dims(
            eval_dict[input_data_fields.groundtruth_instance_masks], axis=0),
        t.uint8)
  images_with_detections = draw_bounding_boxes_on_image_tensors(
      eval_dict[input_data_fields.original_image],
      t.expand_dims(eval_dict[detection_fields.detection_boxes], axis=0),
      t.expand_dims(eval_dict[detection_fields.detection_classes], axis=0),
      t.expand_dims(eval_dict[detection_fields.detection_scores], axis=0),
      category_index,
      instance_masks=instance_masks,
      keypoints=keypoints,
      max_boxes_to_draw=max_boxes_to_draw,
      min_score_thresh=min_score_thresh,
      use_normalized_coordinates=use_normalized_coordinates)
  images_with_groundtruth = draw_bounding_boxes_on_image_tensors(
      eval_dict[input_data_fields.original_image],
      t.expand_dims(eval_dict[input_data_fields.groundtruth_boxes], axis=0),
      t.expand_dims(eval_dict[input_data_fields.groundtruth_classes], axis=0),
      t.expand_dims(
          t.ones_like(
              eval_dict[input_data_fields.groundtruth_classes],
              dtype=t.float32),
          axis=0),
      category_index,
      instance_masks=groundtruth_instance_masks,
      keypoints=None,
      max_boxes_to_draw=None,
      min_score_thresh=0.0,
      use_normalized_coordinates=use_normalized_coordinates)
  return t.concat([images_with_detections, images_with_groundtruth], axis=2)
//  [python]
def draw_side_by_side_evaluation_image(eval_dict,
                                       category_index,
                                       max_boxes_to_draw=20,
                                       min_score_thresh=0.2,
                                       use_normalized_coordinates=True):
  """Creates a side-by-side image with detections and groundtruth.

  Bounding boxes (and instance masks, if available) are visualized on both
  subimages.

  Args:
    eval_dict: The evaluation dictionary returned by
      eval_util.result_dict_for_single_example().
    category_index: A category index (dictionary) produced from a labelmap.
    max_boxes_to_draw: The maximum number of boxes to draw for detections.
    min_score_thresh: The minimum score threshold for showing detections.
    use_normalized_coordinates: Whether to assume boxes and kepoints are in
      normalized coordinates (as opposed to absolute coordiantes).
      Default is True.

  Returns:
    A [1, H, 2 * W, C] uint8 tensor. The subimage on the left corresponds to
      detections, while the subimage on the right corresponds to groundtruth.
  """
  detection_fields = fields.DetectionResultFields()
  input_data_fields = fields.InputDataFields()
  instance_masks = None
  if detection_fields.detection_masks in eval_dict:
    instance_masks = tf.cast(
        tf.expand_dims(eval_dict[detection_fields.detection_masks], axis=0),
        tf.uint8)
  keypoints = None
  if detection_fields.detection_keypoints in eval_dict:
    keypoints = tf.expand_dims(
        eval_dict[detection_fields.detection_keypoints], axis=0)
  groundtruth_instance_masks = None
  if input_data_fields.groundtruth_instance_masks in eval_dict:
    groundtruth_instance_masks = tf.cast(
        tf.expand_dims(
            eval_dict[input_data_fields.groundtruth_instance_masks], axis=0),
        tf.uint8)
  images_with_detections = draw_bounding_boxes_on_image_tensors(
      eval_dict[input_data_fields.original_image],
      tf.expand_dims(eval_dict[detection_fields.detection_boxes], axis=0),
      tf.expand_dims(eval_dict[detection_fields.detection_classes], axis=0),
      tf.expand_dims(eval_dict[detection_fields.detection_scores], axis=0),
      category_index,
      instance_masks=instance_masks,
      keypoints=keypoints,
      max_boxes_to_draw=max_boxes_to_draw,
      min_score_thresh=min_score_thresh,
      use_normalized_coordinates=use_normalized_coordinates)
  images_with_groundtruth = draw_bounding_boxes_on_image_tensors(
      eval_dict[input_data_fields.original_image],
      tf.expand_dims(eval_dict[input_data_fields.groundtruth_boxes], axis=0),
      tf.expand_dims(eval_dict[input_data_fields.groundtruth_classes], axis=0),
      tf.expand_dims(
          tf.ones_like(
              eval_dict[input_data_fields.groundtruth_classes],
              dtype=tf.float32),
          axis=0),
      category_index,
      instance_masks=groundtruth_instance_masks,
      keypoints=None,
      max_boxes_to_draw=None,
      min_score_thresh=0.0,
      use_normalized_coordinates=use_normalized_coordinates)
  return tf.concat([images_with_detections, images_with_groundtruth], axis=2)
//  [python]
def _sample_next_states(self, states, discount = 0.99):
    """Given a sequence of states, samples the `next_states` for loss computation."""
    batch_size = tf.shape(states)[0]
    d = self.sequence_length - 1

    probs = discount ** tf.range(d, dtype=tf.float32)
    probs *= tf.constant([1.0] * (d - 1) + [1.0 / (1 - discount)],
                         dtype=tf.float32)
    probs /= tf.reduce_sum(probs)
    index_dist = tfp.distributions.Categorical(probs=probs, dtype=tf.int64)
    indices = index_dist.sample(batch_size)
    batch_size = tf.to_int64(batch_size)
    next_indices = tf.concat(
        [tf.range(batch_size, dtype=tf.int64)[:, None],
         1 + indices[:, None]], -1)
    next_states = tf.gather_nd(states, next_indices)
    return next_states
//  [python]
def _sample_next_states(self, states, discount = 0.99):
    """Given a sequence of states, samples the `next_states` for loss computation."""
    batch_size = tf.shape(states)[0]
    d = self.sequence_length - 1

    probs = discount ** tf.range(d, dtype=tf.float32)
    probs *= tf.constant([1.0] * (d - 1) + [1.0 / (1 - discount)],
                         dtype=tf.float32)
    probs /= tf.reduce_sum(probs)
    index_dist = tfp.distributions.Categorical(probs=probs, dtype=tf.int64)
    indices = index_dist.sample(batch_size)
    batch_size = tf.cast(batch_size, tf.int64)
    next_indices = tf.concat(
        [tf.range(batch_size, dtype=tf.int64)[:, None],
         1 + indices[:, None]], -1)
    next_states = tf.gather_nd(states, next_indices)
    return next_states
//  [python]
def _prepare_input(self, x):
    """Prepares input for BERT training."""
    batch_size = tf.shape(x)[0]

    # Deal with probability of overlap in masks.
    keep_probability = self.keep_probability
    switch_probability = self.switch_probability / (1 - keep_probability)
    drop_probability = self.drop_probability / (1 - switch_probability) / (1 - keep_probability)
    drop = tf.random.uniform([batch_size, self.attention_length
                             ]) < drop_probability
    switch = tf.random.uniform([batch_size, self.attention_length
                               ]) < switch_probability
    keep = tf.random.uniform([batch_size, self.attention_length
                             ]) < keep_probability

    drop_mask = tf.to_float(drop)
    masked_x = x * (1.0 - drop_mask)[Ellipsis, None] + self.missing_x * drop_mask[Ellipsis, None]

    shuffled_x = tf.concat([masked_x[1:], masked_x[:1]], 0)
    switch_mask = tf.to_float(switch)
    masked_x = masked_x * (1.0 - switch_mask)[Ellipsis, None] + shuffled_x * switch_mask[Ellipsis, None]

    time_encoding = tf.eye(self.attention_length, batch_shape=(batch_size,))
    prepared_x = tf.concat([time_encoding, masked_x], -1)

    full_mask = tf.to_float(drop | switch | keep)

    return prepared_x, full_mask
//  [python]
def _prepare_input(self, x):
    """Prepares input for BERT training."""
    batch_size = tf.shape(x)[0]

    # Deal with probability of overlap in masks.
    keep_probability = self.keep_probability
    switch_probability = self.switch_probability / (1 - keep_probability)
    drop_probability = self.drop_probability / (1 - switch_probability) / (1 - keep_probability)
    drop = tf.random.uniform([batch_size, self.attention_length
                             ]) < drop_probability
    switch = tf.random.uniform([batch_size, self.attention_length
                               ]) < switch_probability
    keep = tf.random.uniform([batch_size, self.attention_length
                             ]) < keep_probability

    drop_mask = tf.cast(drop, tf.float32)
    masked_x = x * (1.0 - drop_mask)[Ellipsis, None] + self.missing_x * drop_mask[Ellipsis, None]

    shuffled_x = tf.concat([masked_x[1:], masked_x[:1]], 0)
    switch_mask = tf.cast(switch, tf.float32)
    masked_x = masked_x * (1.0 - switch_mask)[Ellipsis, None] + shuffled_x * switch_mask[Ellipsis, None]

    time_encoding = tf.eye(self.attention_length, batch_shape=(batch_size,))
    prepared_x = tf.concat([time_encoding, masked_x], -1)

    full_mask = tf.cast(drop | switch | keep, tf.float32)

    return prepared_x, full_mask
//  [python]
def _prepare_input(self, x):
    """Prepares input for BERT training."""
    batch_size = tf.shape(x)[0]

    # Deal with probability of overlap in masks.
    keep_probability = self.keep_probability
    switch_probability = self.switch_probability / (1 - keep_probability)
    drop_probability = self.drop_probability / (1 - switch_probability) / (1 - keep_probability)
    drop = tf.random.uniform(tf.shape(x)) < drop_probability
    switch = tf.random.uniform(tf.shape(x)) < switch_probability
    keep = tf.random.uniform(tf.shape(x)) < keep_probability

    drop_mask = tf.to_float(drop)
    masked_x = x * (1.0 - drop_mask)

    shuffled_x = tf.concat([masked_x[1:], masked_x[:1]], 0)
    switch_mask = tf.to_float(switch)
    masked_x = masked_x * (1.0 - switch_mask) + shuffled_x * switch_mask

    full_mask = tf.to_float(drop | switch | keep)

    return masked_x, full_mask, drop
//  [python]
def _prepare_input(self, x):
    """Prepares input for BERT training."""
    batch_size = tf.shape(x)[0]

    # Deal with probability of overlap in masks.
    keep_probability = self.keep_probability
    switch_probability = self.switch_probability / (1 - keep_probability)
    drop_probability = self.drop_probability / (1 - switch_probability) / (1 - keep_probability)
    drop = tf.random.uniform(tf.shape(x)) < drop_probability
    switch = tf.random.uniform(tf.shape(x)) < switch_probability
    keep = tf.random.uniform(tf.shape(x)) < keep_probability

    drop_mask = tf.cast(drop, tf.float32)
    masked_x = x * (1.0 - drop_mask)

    shuffled_x = tf.concat([masked_x[1:], masked_x[:1]], 0)
    switch_mask = tf.cast(switch, tf.float32)
    masked_x = masked_x * (1.0 - switch_mask) + shuffled_x * switch_mask

    full_mask = tf.cast(drop | switch | keep, tf.float32)

    return masked_x, full_mask, drop
//  [python]
def fit(self, states, actions):
    """Updates critic parameters.

    Args:
      states: Batch of sequences of states.
      actions: Batch of sequences of actions.

    Returns:
      Dictionary with information to track.
    """
    batch_size = tf.shape(states)[0]
    with tf.GradientTape(watch_accessed_variables=False) as tape:
      tape.watch(self.all_variables)

      next_states = self._sample_next_states(states)
      cur_states = states[:, 0, :]

      cur_embed = self.embedder(cur_states, stop_gradient=False)
      next_embed = self.embedder(next_states, stop_gradient=False)

      # Update history of embeddings with this batch's next_embed.
      self.embed_history.assign(tf.concat([self.embed_history[batch_size:], next_embed], 0))

      action_embed_input = tf.concat(
          [cur_states, tf.reshape(actions[:, :-1, :], [batch_size, -1])], -1)
      action_embed = self.action_embedder(action_embed_input, stop_gradient=False)

      tau = 2.0
      energy_fn = lambda z: -tau * tf.reduce_sum(huber(z), -1)
      positive_loss = tf.reduce_mean(energy_fn(cur_embed + action_embed - next_embed))
      # Negative loss should just be a log-avg-exp, but we compute it in a more
      # numerically stable way below.
      prior_log_probs = tf.reduce_logsumexp(
          energy_fn((cur_embed + action_embed)[:, None, :]
                    - self.embed_history[None, :, :]),
          axis=-1) - tf.math.log(tf.cast(self.embed_history.shape[0], tf.float32))
      shifted_next_embed = tf.concat([next_embed[1:], next_embed[:1]], 0)
      negative_loss = tf.reduce_mean(
          tf.exp(energy_fn((cur_embed + action_embed) - shifted_next_embed)
                 - tf.stop_gradient(prior_log_probs)))

      loss = tf.reduce_mean(-positive_loss + negative_loss)

    grads = tape.gradient(loss, self.all_variables)

    self.optimizer.apply_gradients(
        zip(grads, self.all_variables))

    return {
        'embed_loss': loss
    }
//  [python]
def _prepare_input(self, x, missing_x):
    """Prepares input for BERT training."""
    batch_size = tf.shape(x)[0]

    # Deal with probability of overlap in masks.
    keep_probability = self.keep_probability
    switch_probability = self.switch_probability / (1 - keep_probability)
    drop_probability = self.drop_probability / (1 - switch_probability) / (
        1 - keep_probability)
    drop = tf.random.uniform([batch_size, self.attention_length
                             ]) < drop_probability
    switch = tf.random.uniform([batch_size, self.attention_length
                               ]) < switch_probability
    keep = tf.random.uniform([batch_size, self.attention_length
                             ]) < keep_probability

    drop_mask = tf.to_float(drop)
    masked_x = x * (1.0 - drop_mask)[Ellipsis, None] + missing_x * drop_mask[Ellipsis,
                                                                        None]

    shuffled_x = tf.concat([masked_x[1:], masked_x[:1]], 0)
    switch_mask = tf.to_float(switch)
    masked_x = masked_x * (
        1.0 - switch_mask)[Ellipsis, None] + shuffled_x * switch_mask[Ellipsis, None]

    full_mask = tf.to_float(drop | switch | keep)

    return masked_x, full_mask
//  [python]
def _prepare_input(self, x, missing_x):
    """Prepares input for BERT training."""
    batch_size = tf.shape(x)[0]

    # Deal with probability of overlap in masks.
    keep_probability = self.keep_probability
    switch_probability = self.switch_probability / (1 - keep_probability)
    drop_probability = self.drop_probability / (1 - switch_probability) / (
        1 - keep_probability)
    drop = tf.random.uniform([batch_size, self.attention_length
                             ]) < drop_probability
    switch = tf.random.uniform([batch_size, self.attention_length
                               ]) < switch_probability
    keep = tf.random.uniform([batch_size, self.attention_length
                             ]) < keep_probability

    drop_mask = tf.cast(drop, tf.float32)
    masked_x = x * (1.0 - drop_mask)[Ellipsis, None] + missing_x * drop_mask[Ellipsis,
                                                                        None]

    shuffled_x = tf.concat([masked_x[1:], masked_x[:1]], 0)
    switch_mask = tf.cast(switch, tf.float32)
    masked_x = masked_x * (
        1.0 - switch_mask)[Ellipsis, None] + shuffled_x * switch_mask[Ellipsis, None]

    full_mask = tf.cast(drop | switch | keep, tf.float32)

    return masked_x, full_mask
//  [python]
def augment_bounding_boxes(
        self, bounding_boxes, transformation, image, **kwargs
    ):
        if self.bounding_box_format is None:
            raise ValueError(
                "Please provide a `bounding_box_format` when augmenting "
                "bounding boxes with `RandomAspectRatio()`."
            )
        bounding_boxes = bounding_boxes.copy()
        img_shape = tf.shape(image)
        img_shape = t.cast(img_shape, self.compute_dtype)
        height, width = img_shape[0], img_shape[1]
        height = height / transformation
        width = width * transformation

        bounding_boxes = bounding_box.convert_format(
            bounding_boxes,
            source=self.bounding_box_format,
            target="xyxy",
            image_shape=img_shape,
        )
        x, y, x2, y2 = tf.split(bounding_boxes["boxes"], [1, 1, 1, 1], axis=-1)
        x = x * transformation
        x2 = x2 * transformation
        y = y / transformation
        y2 = y2 / transformation
        boxes = tf.concat([x, y, x2, y2], axis=-1)
        boxes = bounding_box.convert_format(
            boxes,
            source="xyxy",
            target=self.bounding_box_format,
            image_shape=tf.stack([height, width, 3], axis=0),
        )
        bounding_boxes["boxes"] = boxes
        return bounding_boxes
//  [python]
def augment_bounding_boxes(
        self, bounding_boxes, transformation, image, **kwargs
    ):
        if self.bounding_box_format is None:
            raise ValueError(
                "Please provide a `bounding_box_format` when augmenting "
                "bounding boxes with `RandomAspectRatio()`."
            )
        bounding_boxes = bounding_boxes.copy()
        img_shape = tf.shape(image)
        img_shape = tf.cast(img_shape, self.compute_dtype)
        height, width = img_shape[0], img_shape[1]
        height = height / transformation
        width = width * transformation

        bounding_boxes = bounding_box.convert_format(
            bounding_boxes,
            source=self.bounding_box_format,
            target="xyxy",
            image_shape=img_shape,
        )
        x, y, x2, y2 = tf.split(bounding_boxes["boxes"], [1, 1, 1, 1], axis=-1)
        x = x * transformation
        x2 = x2 * transformation
        y = y / transformation
        y2 = y2 / transformation
        boxes = tf.concat([x, y, x2, y2], axis=-1)
        boxes = bounding_box.convert_format(
            boxes,
            source="xyxy",
            target=self.bounding_box_format,
            image_shape=tf.stack([height, width, 3], axis=0),
        )
        bounding_boxes["boxes"] = boxes
        return bounding_boxes
//  [python]
def augment_image(self, image, transformation, **kwargs):
        # images....transformation
        img_shape = tf.to_int32(tf.shape(image))
        height, width = img_shape[0], img_shape[1]
        height = height / transformation
        width = width * transformation

        target_size = tf.to_int32(tf.stack([height, width]))
        result = tf.image.resize(
            image, size=target_size, method=self.interpolation
        )
        return tf.to_float(result)
//  [python]
def augment_image(self, image, transformation, **kwargs):
        # images....transformation
        img_shape = tf.cast(tf.shape(image), self.compute_dtype)
        height, width = img_shape[0], img_shape[1]
        height = height / transformation
        width = width * transformation

        target_size = tf.cast(tf.stack([height, width]), tf.int32)
        result = tf.image.resize(
            image, size=target_size, method=self.interpolation
        )
        return tf.cast(result, self.compute_dtype)
//  [python]
def fit(self, states, actions):
    """Updates critic parameters.

    Args:
      states: Batch of sequences of states.
      actions: Batch of sequences of actions.

    Returns:
      Dictionary with information to track.
    """
    states = states[:, :self.attention_length, :]
    actions = actions[:, :self.attention_length, :]
    batch_size = tf.shape(states)[0]
    with tf.GradientTape(watch_accessed_variables=False) as tape:
      tape.watch(self.all_variables)

      all_states = tf.reshape(states, [batch_size * self.attention_length, -1])
      all_embeddings = self.embedder(all_states, stop_gradient=False)
      embeddings = tf.reshape(
          all_embeddings,
          [batch_size, self.attention_length, self.embedding_dim])

      states_in = embeddings if self.embed_on_input else states
      actions_in = actions
      if self.input_dimension_dropout > 0:
        states_in *= tf.to_float(
            tf.random.uniform(tf.shape(states_in)) < self.input_dimension_dropout)
        states_in *= 1 / (1 - self.input_dimension_dropout)
        actions_in *= tf.to_float(
            tf.random.uniform(tf.shape(actions_in)) < self.input_dimension_dropout)
        actions_in *= 1 / (1 - self.input_dimension_dropout)
      if self.predict_actions:
        states_in, drop_mask = self._prepare_input(states_in)
        attention_in = states_in
      elif self.modify_actions:
        attention_in = tf.concat([states_in, actions_in], -1)
        attention_in, drop_mask = self._prepare_input(attention_in)
      else:
        states_in, drop_mask = self._prepare_input(states_in)
        attention_in = tf.concat([states_in, actions_in], -1)
      attention_out = self.transformer(attention_in, training=True)

      mask_indices = tf.where(drop_mask > 0.0)
      if self.predict_actions:
        pred_actions = tf.gather_nd(attention_out, mask_indices)
        true_actions = tf.gather_nd(actions, mask_indices)
        loss = tf.reduce_mean(tf.reduce_sum((pred_actions - true_actions) ** 2, -1))
      else:
        pred_embeddings = tf.gather_nd(attention_out, mask_indices)
        true_embeddings = tf.gather_nd(embeddings, mask_indices)

        energies = self.compute_energy(pred_embeddings, true_embeddings)
        positive_loss = tf.linalg.diag_part(energies)
        negative_loss = tf.reduce_logsumexp(energies, axis=-1)

        loss = tf.reduce_mean(-positive_loss + negative_loss)

    grads = tape.gradient(loss, self.all_variables)

    self.optimizer.apply_gradients(
        zip(grads, self.all_variables))

    return {
        'embed_loss': loss
    }
//  [python]
def fit(self, states, actions):
    """Updates critic parameters.

    Args:
      states: Batch of sequences of states.
      actions: Batch of sequences of actions.

    Returns:
      Dictionary with information to track.
    """
    states = states[:, :self.attention_length, :]
    actions = actions[:, :self.attention_length, :]
    batch_size = tf.shape(states)[0]
    with tf.GradientTape(watch_accessed_variables=False) as tape:
      tape.watch(self.all_variables)

      all_states = tf.reshape(states, [batch_size * self.attention_length, -1])
      all_embeddings = self.embedder(all_states, stop_gradient=False)
      embeddings = tf.reshape(
          all_embeddings,
          [batch_size, self.attention_length, self.embedding_dim])

      states_in = embeddings if self.embed_on_input else states
      actions_in = actions
      if self.input_dimension_dropout > 0:
        states_in *= tf.cast(
            tf.random.uniform(tf.shape(states_in)) < self.input_dimension_dropout,
            tf.float32)
        states_in *= 1 / (1 - self.input_dimension_dropout)
        actions_in *= tf.cast(
            tf.random.uniform(tf.shape(actions_in)) < self.input_dimension_dropout,
            tf.float32)
        actions_in *= 1 / (1 - self.input_dimension_dropout)
      if self.predict_actions:
        states_in, drop_mask = self._prepare_input(states_in)
        attention_in = states_in
      elif self.modify_actions:
        attention_in = tf.concat([states_in, actions_in], -1)
        attention_in, drop_mask = self._prepare_input(attention_in)
      else:
        states_in, drop_mask = self._prepare_input(states_in)
        attention_in = tf.concat([states_in, actions_in], -1)
      attention_out = self.transformer(attention_in, training=True)

      mask_indices = tf.where(drop_mask > 0.0)
      if self.predict_actions:
        pred_actions = tf.gather_nd(attention_out, mask_indices)
        true_actions = tf.gather_nd(actions, mask_indices)
        loss = tf.reduce_mean(tf.reduce_sum((pred_actions - true_actions) ** 2, -1))
      else:
        pred_embeddings = tf.gather_nd(attention_out, mask_indices)
        true_embeddings = tf.gather_nd(embeddings, mask_indices)

        energies = self.compute_energy(pred_embeddings, true_embeddings)
        positive_loss = tf.linalg.diag_part(energies)
        negative_loss = tf.reduce_logsumexp(energies, axis=-1)

        loss = tf.reduce_mean(-positive_loss + negative_loss)

    grads = tape.gradient(loss, self.all_variables)

    self.optimizer.apply_gradients(
        zip(grads, self.all_variables))

    return {
        'embed_loss': loss
    }
//  [python]
def get_sampling_offset(sequence: tf.Tensor,
                        num_steps: Optional[int],
                        is_training: bool,
                        stride: int = 1,
                        seed: Optional[int] = None) -> tf.Tensor:
  """Calculates the initial offset for a sequence where all steps will fit.

  Args:
    sequence: any tensor where the first dimension is timesteps.
    num_steps: The number of timesteps we will output. If None,
      deterministically start at the first frame.
    is_training: A boolean indicates whether the graph is for training or not.
      If False, the starting frame always the first frame.
    stride: distance to sample between timesteps.
    seed: a deterministic seed to use when sampling.
  Returns:
    The first index to begin sampling from. A best effort is made to provide a
    starting index such that all requested steps fit within the sequence (i.e.
    offset + 1 + (num_steps - 1) * stride < len(sequence)). If this is not
    satisfied, the starting index is chosen randomly from the full sequence.
  """
  if num_steps is None or not is_training:
    return tf.constant(0)
  sequence_length = tf.shape(sequence)[0]
  max_offset = tf.cond(
      tf.greater(sequence_length, (num_steps - 1) * stride),
      lambda: sequence_length - (num_steps - 1) * stride,
      lambda: sequence_length)
  offset = tf.random.uniform(
      (),
      maxval=t.cast(max_offset, tf.int32),
      dtype=tf.int32,
      seed=seed)
  return offset
//  [python]
def get_sampling_offset(sequence: tf.Tensor,
                        num_steps: Optional[int],
                        is_training: bool,
                        stride: int = 1,
                        seed: Optional[int] = None) -> tf.Tensor:
  """Calculates the initial offset for a sequence where all steps will fit.

  Args:
    sequence: any tensor where the first dimension is timesteps.
    num_steps: The number of timesteps we will output. If None,
      deterministically start at the first frame.
    is_training: A boolean indicates whether the graph is for training or not.
      If False, the starting frame always the first frame.
    stride: distance to sample between timesteps.
    seed: a deterministic seed to use when sampling.
  Returns:
    The first index to begin sampling from. A best effort is made to provide a
    starting index such that all requested steps fit within the sequence (i.e.
    offset + 1 + (num_steps - 1) * stride < len(sequence)). If this is not
    satisfied, the starting index is chosen randomly from the full sequence.
  """
  if num_steps is None or not is_training:
    return tf.constant(0)
  sequence_length = tf.shape(sequence)[0]
  max_offset = tf.cond(
      tf.greater(sequence_length, (num_steps - 1) * stride),
      lambda: sequence_length - (num_steps - 1) * stride,
      lambda: sequence_length)
  offset = tf.random.uniform(
      (),
      maxval=tf.cast(max_offset, tf.int32),
      dtype=tf.int32,
      seed=seed)
  return offset
//  [python]
def resize_smallest(frames: tf.Tensor, min_resize: int) -> tf.Tensor:
  """Resizes frames so that min(height, width) is equal to min_resize.

  This function will not do anything if the min(height, width) is already equal
  to min_resize. This allows to save compute time.

  Args:
    frames: A Tensor of dimension [timesteps, input_h, input_w, channels].
    min_resize: Minimum size of the final image dimensions.
  Returns:
    A Tensor of shape [timesteps, output_h, output_w, channels] of type
      frames.dtype where min(output_h, output_w) = min_resize.
  """
  shape = tf.shape(frames)
  input_h = shape[1]
  input_w = shape[2]

  output_h = tf.maximum(min_resize, (input_h * min_resize) // input_w)
  output_w = tf.maximum(min_resize, (input_w * min_resize) // input_h)

  def resize_fn():
    frames_resized = tf.image.resize(frames, (output_h, output_w))
    return tf.to_float(frames_resized)

  should_resize = tf.math.logical_or(tf.not_equal(input_w, output_w),
                                     tf.not_equal(input_h, output_h))
  frames = tf.cond(should_resize, resize_fn, lambda: frames)

  return frames
//  [python]
def resize_smallest(frames: tf.Tensor, min_resize: int) -> tf.Tensor:
  """Resizes frames so that min(height, width) is equal to min_resize.

  This function will not do anything if the min(height, width) is already equal
  to min_resize. This allows to save compute time.

  Args:
    frames: A Tensor of dimension [timesteps, input_h, input_w, channels].
    min_resize: Minimum size of the final image dimensions.
  Returns:
    A Tensor of shape [timesteps, output_h, output_w, channels] of type
      frames.dtype where min(output_h, output_w) = min_resize.
  """
  shape = tf.shape(frames)
  input_h = shape[1]
  input_w = shape[2]

  output_h = tf.maximum(min_resize, (input_h * min_resize) // input_w)
  output_w = tf.maximum(min_resize, (input_w * min_resize) // input_h)

  def resize_fn():
    frames_resized = tf.image.resize(frames, (output_h, output_w))
    return tf.cast(frames_resized, frames.dtype)

  should_resize = tf.math.logical_or(tf.not_equal(input_w, output_w),
                                     tf.not_equal(input_h, output_h))
  frames = tf.cond(should_resize, resize_fn, lambda: frames)

  return frames
//  [python]
def process_samples(features_dict, num_frames=32, stride=1, is_training=True,
                    min_resize=224, crop_size=224, num_windows=1):
  """Process video frames."""

  video = features_dict['video']

  if is_training:
    assert num_windows == 1
    video = random_sample_sequence(video, num_frames, stride)
    is_flipped = tf.random.uniform((), minval=0, maxval=2, dtype=tf.int32)
    video = tf.cond(tf.equal(is_flipped, 1),
                    true_fn=lambda: tf.image.flip_left_right(video),
                    false_fn=lambda: video)
  else:
    video = sample_linspace_sequence(video, num_windows, num_frames, stride)

  # Resize smallest side.
  video = resize_smallest(video, min_resize)

  if is_training:
    # Random crop.
    video = tf.image.random_crop(video, [num_frames, crop_size, crop_size, 3])
  else:
    # Central crop.
    video = tf.image.resize_with_crop_or_pad(video, crop_size, crop_size)

  video = tf.to_float(video)
  video /= 255.0  # Set between [0, 1].

  features_dict['video'] = video
  return features_dict
//  [python]
def process_samples(features_dict, num_frames=32, stride=1, is_training=True,
                    min_resize=224, crop_size=224, num_windows=1):
  """Process video frames."""

  video = features_dict['video']

  if is_training:
    assert num_windows == 1
    video = random_sample_sequence(video, num_frames, stride)
    is_flipped = tf.random.uniform((), minval=0, maxval=2, dtype=tf.int32)
    video = tf.cond(tf.equal(is_flipped, 1),
                    true_fn=lambda: tf.image.flip_left_right(video),
                    false_fn=lambda: video)
  else:
    video = sample_linspace_sequence(video, num_windows, num_frames, stride)

  # Resize smallest side.
  video = resize_smallest(video, min_resize)

  if is_training:
    # Random crop.
    video = tf.image.random_crop(video, [num_frames, crop_size, crop_size, 3])
  else:
    # Central crop.
    video = tf.image.resize_with_crop_or_pad(video, crop_size, crop_size)

  video = tf.cast(video, tf.float32)
  video /= 255.0  # Set between [0, 1].

  features_dict['video'] = video
  return features_dict
//  [python]
def sample_linspace_sequence(sequence: tf.Tensor,
                             num_windows: int,
                             num_steps: int,
                             stride: int = 1) -> tf.Tensor:
  """Samples num_windows segments from sequence with linearly spaced offsets.

  The samples are concatenated in a single Tensor in order to have the same
  format structure per timestep (e.g. a single frame). If num_steps * stride is
  bigger than the number of timesteps, the sequence is repeated. This function
  can be used in evaluation in order to extract enough segments in order to span
  the entire sequence.

  Args:
    sequence: Any tensor where the first dimension is timesteps.
    num_windows: Number of windows retrieved from the sequence.
    num_steps: Number of steps (e.g. frames) to take.
    stride: Distance to sample between timesteps.

  Returns:
    A single Tensor with first dimension num_windows * num_steps. The Tensor
    contains the concatenated list of num_windows tensors which offsets have
    been linearly spaced from input.
  """
  sequence_length = tf.shape(sequence)[0]
  max_offset = tf.maximum(0, sequence_length - num_steps * stride)
  offsets = tf.linspace(0.0, tf.cast(max_offset, tf.float32), num_windows)
  offsets = tf.to_int32(offsets)

  all_indices = []
  for i in range(num_windows):
    all_indices.append(
        sample_or_pad_sequence_indices(
            sequence=sequence,
            num_steps=num_steps,
            is_training=False,
            repeat_sequence=True,  # Will repeat the sequence if request more.
            stride=stride,
            offset=offsets[i]))

  indices = tf.concat(all_indices, axis=0)
  indices.set_shape((num_windows * num_steps,))
  output = tf.gather(sequence, indices)

  return output
//  [python]
def sample_linspace_sequence(sequence: tf.Tensor,
                             num_windows: int,
                             num_steps: int,
                             stride: int = 1) -> tf.Tensor:
  """Samples num_windows segments from sequence with linearly spaced offsets.

  The samples are concatenated in a single Tensor in order to have the same
  format structure per timestep (e.g. a single frame). If num_steps * stride is
  bigger than the number of timesteps, the sequence is repeated. This function
  can be used in evaluation in order to extract enough segments in order to span
  the entire sequence.

  Args:
    sequence: Any tensor where the first dimension is timesteps.
    num_windows: Number of windows retrieved from the sequence.
    num_steps: Number of steps (e.g. frames) to take.
    stride: Distance to sample between timesteps.

  Returns:
    A single Tensor with first dimension num_windows * num_steps. The Tensor
    contains the concatenated list of num_windows tensors which offsets have
    been linearly spaced from input.
  """
  sequence_length = tf.shape(sequence)[0]
  max_offset = tf.maximum(0, sequence_length - num_steps * stride)
  offsets = tf.linspace(0.0, tf.cast(max_offset, tf.float32), num_windows)
  offsets = tf.cast(offsets, tf.int32)

  all_indices = []
  for i in range(num_windows):
    all_indices.append(
        sample_or_pad_sequence_indices(
            sequence=sequence,
            num_steps=num_steps,
            is_training=False,
            repeat_sequence=True,  # Will repeat the sequence if request more.
            stride=stride,
            offset=offsets[i]))

  indices = tf.concat(all_indices, axis=0)
  indices.set_shape((num_windows * num_steps,))
  output = tf.gather(sequence, indices)

  return output
//  [python]
def unpack_labels(self, labels,
                    is_box = False):
    """Unpacks an array of labels into multiscales labels.

    Args:
      labels: labels to unpack.
      is_box: to unpack anchor boxes or not. If it is true, will unpack to 2D,
        otherwise, will unpack to 3D.

    Returns:
      unpacked_labels: a dictionary contains unpack labels in different levels.
    """
    unpacked_labels = {}
    count = 0
    for level in range(self.min_level, self.max_level + 1):
      feat_size_y = t.cast(self.image_size[0] / 2 ** level, t.int32)
      feat_size_x = t.cast(self.image_size[1] / 2 ** level, t.int32)
      steps = feat_size_y * feat_size_x * self.anchors_per_location
      if is_box:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [-1, 4])
      else:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [feat_size_y, feat_size_x, -1])
      count += steps
    return unpacked_labels
//  [python]
def unpack_labels(self, labels,
                    is_box = False):
    """Unpacks an array of labels into multiscales labels.

    Args:
      labels: labels to unpack.
      is_box: to unpack anchor boxes or not. If it is true, will unpack to 2D,
        otherwise, will unpack to 3D.

    Returns:
      unpacked_labels: a dictionary contains unpack labels in different levels.
    """
    unpacked_labels = {}
    count = 0
    for level in range(self.min_level, self.max_level + 1):
      feat_size_y = tf.cast(self.image_size[0] / 2 ** level, tf.int32)
      feat_size_x = tf.cast(self.image_size[1] / 2 ** level, tf.int32)
      steps = feat_size_y * feat_size_x * self.anchors_per_location
      if is_box:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [-1, 4])
      else:
        unpacked_labels[level] = tf.reshape(labels[count:count + steps],
                                            [feat_size_y, feat_size_x, -1])
      count += steps
    return unpacked_labels
//  [python]
def decode_example(
    serialized_proto: str, features_and_vocabs_by_name: FeaturesAndVocabsByName,
    label_feature_name: str) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
  """Decode single serialized example.

  Decode single serialized example, accoring to specified features in input
  config. Perform vocabulary lookup if vocabulary is specified.

  Args:
    serialized_proto: The serialized proto that needs to be decoded.
    features_and_vocabs_by_name: A FeaturesAndVocabsByName object containing
      features and vocabs dictionaries by feature names.
    label_feature_name: Name of the label feature.

  Returns:
    features: The decoded features dictionary.
    label_feature: The label feature.

  """
  features_by_name = features_and_vocabs_by_name.features_by_name
  vocabs_by_name = features_and_vocabs_by_name.vocabs_by_name
  name_to_features = {}
  name_to_features[label_feature_name] = tf.io.FixedLenFeature([1], tf.int64)
  for feature_name, feature in features_by_name.items():
    name_to_features[feature_name] = _get_feature_spec(
        feature_type=feature.feature_type,
        feature_length=feature.feature_length)
  record_features = tf.io.parse_single_example(serialized_proto,
                                               name_to_features)

  features = {}
  for feature_name, feature_value in record_features.items():
    if feature_name in vocabs_by_name:
      features[feature_name] = vocabs_by_name[feature_name].lookup(
          feature_value)
    else:
      features[feature_name] = feature_value
  features = {
      k: v if v.dtype != tf.int64 else tf.to_int32(v)
      for k, v in features.items()
  }
  label_feature = features[label_feature_name]
  return features, label_feature
//  [python]
def decode_example(
    serialized_proto: str, features_and_vocabs_by_name: FeaturesAndVocabsByName,
    label_feature_name: str) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
  """Decode single serialized example.

  Decode single serialized example, accoring to specified features in input
  config. Perform vocabulary lookup if vocabulary is specified.

  Args:
    serialized_proto: The serialized proto that needs to be decoded.
    features_and_vocabs_by_name: A FeaturesAndVocabsByName object containing
      features and vocabs dictionaries by feature names.
    label_feature_name: Name of the label feature.

  Returns:
    features: The decoded features dictionary.
    label_feature: The label feature.

  """
  features_by_name = features_and_vocabs_by_name.features_by_name
  vocabs_by_name = features_and_vocabs_by_name.vocabs_by_name
  name_to_features = {}
  name_to_features[label_feature_name] = tf.io.FixedLenFeature([1], tf.int64)
  for feature_name, feature in features_by_name.items():
    name_to_features[feature_name] = _get_feature_spec(
        feature_type=feature.feature_type,
        feature_length=feature.feature_length)
  record_features = tf.io.parse_single_example(serialized_proto,
                                               name_to_features)

  features = {}
  for feature_name, feature_value in record_features.items():
    if feature_name in vocabs_by_name:
      features[feature_name] = vocabs_by_name[feature_name].lookup(
          feature_value)
    else:
      features[feature_name] = feature_value
  features = {
      k: v if v.dtype != tf.int64 else tf.cast(v, tf.int32)
      for k, v in features.items()
  }
  label_feature = features[label_feature_name]
  return features, label_feature
//  [python]
def discretize_rtg(self, rtg):
    if self.min_rtg and self.max_rtg:
      rtg = (rtg - self.min_rtg) / (self.max_rtg - self.min_rtg) * self.num_rtg
    return t.cast(rtg, t.int32)
//  [python]
def discretize_rtg(self, rtg):
    if self.min_rtg and self.max_rtg:
      rtg = (rtg - self.min_rtg) / (self.max_rtg - self.min_rtg) * self.num_rtg
    return tf.cast(rtg, tf.int32)
//  [python]
def train_step(self, transitions):
    states, actions, rewards, _, mask = transitions
    # Replace padded -1s with 0s.
    states = tf.where(states < 0, 0, states)
    actions = tf.where(actions < 0, 0, actions)
    rewards = tf.where(rewards < 0., 0., rewards)
    states = tf.squeeze(states)
    actions = tf.squeeze(actions)
    rewards = tf.squeeze(tf.to_int32(rewards))
    rtgs = rewards
    if tf.rank(rtgs) > 1:
      rtgs = tf.math.cumsum(rewards, axis=-1, reverse=True)
    with tf.GradientTape() as tape:
      rtg = self.discretize_rtg(rtgs)
      preds = tf.gather_nd(self.policy, tf.stack([states, rtg], axis=-1))
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(actions, preds)
      loss = tf.reduce_mean(loss)
    grads = tape.gradient(loss, self.variables)
    self.optimizer.apply_gradients(zip(grads, self.variables))
    return {'loss': loss}
//  [python]
def train_step(self, transitions):
    states, actions, rewards, _, mask = transitions
    # Replace padded -1s with 0s.
    states = tf.where(states < 0, 0, states)
    actions = tf.where(actions < 0, 0, actions)
    rewards = tf.where(rewards < 0., 0., rewards)
    states = tf.squeeze(states)
    actions = tf.squeeze(actions)
    rewards = tf.squeeze(tf.cast(rewards, tf.int32))
    rtgs = rewards
    if tf.rank(rtgs) > 1:
      rtgs = tf.math.cumsum(rewards, axis=-1, reverse=True)
    with tf.GradientTape() as tape:
      rtg = self.discretize_rtg(rtgs)
      preds = tf.gather_nd(self.policy, tf.stack([states, rtg], axis=-1))
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(actions, preds)
      loss = tf.reduce_mean(loss)
    grads = tape.gradient(loss, self.variables)
    self.optimizer.apply_gradients(zip(grads, self.variables))
    return {'loss': loss}
//  [python]
def _cast_or_convert(original: tf.Tensor, target_type: tf.DType) -> tf.Tensor:
  if target_type == tf.string and original.dtype != tf.string:
    return tf.as_string(original)
  else:
    return tf.cast(original, target_type)
//  [python]
def _make_transform_fn(
    bert_tokenizer,
    bert_cls_token_id,
    bert_sep_token_id,
):
  """Prepares the transformation function."""
  @tf.function
  def _prepare_for_bert(sample):
    """Prepares a question sample from ELI5 to be fed to BERT."""
    bert_question_token_ids = bert_tokenizer.tokenize(
        tf.expand_dims(sample["question"], 0))
    bert_question_token_ids = t.cast(
        bert_question_token_ids.merge_dims(1, 2).to_tensor(), tf.int32)
    cls_ids = tf.fill([tf.shape(bert_question_token_ids)[0], 1],
                      bert_cls_token_id)
    sep_ids = tf.fill([tf.shape(bert_question_token_ids)[0], 1],
                      bert_sep_token_id)
    bert_question_token_ids = tf.concat(
        (cls_ids, bert_question_token_ids, sep_ids), 1)

    return dict(
        bert_question_token_ids=bert_question_token_ids,
        bert_attention_masks=tf.ones_like(bert_question_token_ids),
        **sample
    )

  return _prepare_for_bert
//  [python]
def _make_transform_fn(
    bert_tokenizer,
    bert_cls_token_id,
    bert_sep_token_id,
):
  """Prepares the transformation function."""
  @tf.function
  def _prepare_for_bert(sample):
    """Prepares a question sample from ELI5 to be fed to BERT."""
    bert_question_token_ids = bert_tokenizer.tokenize(
        tf.expand_dims(sample["question"], 0))
    bert_question_token_ids = tf.cast(
        bert_question_token_ids.merge_dims(1, 2).to_tensor(), tf.int32)
    cls_ids = tf.fill([tf.shape(bert_question_token_ids)[0], 1],
                      bert_cls_token_id)
    sep_ids = tf.fill([tf.shape(bert_question_token_ids)[0], 1],
                      bert_sep_token_id)
    bert_question_token_ids = tf.concat(
        (cls_ids, bert_question_token_ids, sep_ids), 1)

    return dict(
        bert_question_token_ids=bert_question_token_ids,
        bert_attention_masks=tf.ones_like(bert_question_token_ids),
        **sample
    )

  return _prepare_for_bert