jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _parse_train_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for training."""
    # Process image and label.
    image = decoded_tensors[self._image_key]
    image = _process_image(
        image=image,
        is_training=True,
        is_ssl=self._is_ssl,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
        crop_size=self._crop_size,
        zero_centering_image=self._zero_centering_image)
    image = tf.cast(image, dtype=self._dtype)
    features = {'image': image}

    label = decoded_tensors[self._label_key]
    label = _process_label(label, self._one_hot_label, self._num_classes)

    return features, label
//  [python]
def __call__(self, projection1: tf.Tensor, projection2: tf.Tensor):
    """Compute the contrastive loss for contrastive learning.

    Note that projection2 is generated with the same batch (same order) of raw
    images, but with different augmentation. More specifically:
    image[i] -> random augmentation 1 -> projection -> projection1[i]
    image[i] -> random augmentation 2 -> projection -> projection2[i]

    Args:
      projection1: projection vector of shape (bsz, dim).
      projection2: projection vector of shape (bsz, dim).

    Returns:
      A loss scalar.
      The logits for contrastive prediction task.
      The labels for contrastive prediction task.
    """
    # Get (normalized) hidden1 and hidden2.
    if self._projection_norm:
      projection1 = tf.math.l2_normalize(projection1, -1)
      projection2 = tf.math.l2_normalize(projection2, -1)
    batch_size = tf.shape(projection1)[0]

    p1_local, p2_local = projection1, projection2
    # Gather projection1/projection2 across replicas and create local labels.
    num_replicas_in_sync = tf.distribute.get_strategy().num_replicas_in_sync
    if num_replicas_in_sync > 1:
      p1_global = cross_replica_concat(p1_local, num_replicas_in_sync)
      p2_global = cross_replica_concat(p2_local, num_replicas_in_sync)
      global_batch_size = tf.shape(p1_global)[0]

      replica_context = tf.distribute.get_replica_context()
      replica_id = tf.cast(
          tf.cast(replica_context.replica_id_in_sync_group, tf.uint32),
          tf.int32)
      labels_idx = tf.range(batch_size) + replica_id * batch_size
      labels = tf.one_hot(labels_idx, global_batch_size * 2)
      masks = tf.one_hot(labels_idx, global_batch_size)
    else:
      p1_global = p1_local
      p2_global = p2_local
      labels = tf.one_hot(tf.range(batch_size), batch_size * 2)
      masks = tf.one_hot(tf.range(batch_size), batch_size)

    tb_matmul = functools.partial(tf.matmul, transpose_b=True)

    logits_aa = tb_matmul(p1_local, p1_global) / self._temperature
    logits_aa = logits_aa - masks * LARGE_NUM

    logits_bb = tb_matmul(p2_local, p2_global) / self._temperature
    logits_bb = logits_bb - masks * LARGE_NUM

    logits_ab = tb_matmul(p1_local, p2_global) / self._temperature
    logits_ba = tb_matmul(p2_local, p1_global) / self._temperature

    loss_a_local = tf.nn.softmax_cross_entropy_with_logits(
        labels, tf.concat([logits_ab, logits_aa], 1))
    loss_b_local = tf.nn.softmax_cross_entropy_with_logits(
        labels, tf.concat([logits_ba, logits_bb], 1))
    loss_local = tf.reduce_mean(loss_a_local + loss_b_local)

    return loss_local, (logits_ab, labels)
//  [python]
def _parse_eval_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for evaluation."""
    image = decoded_tensors[self._image_key]
    image = _process_image(
        image=image,
        is_training=False,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
        crop_size=self._crop_size,
        num_crops=self._num_crops,
        zero_centering_image=self._zero_centering_image)
    image = t.cast(image, dtype=self._dtype)
    features = {'image': image}

    label = decoded_tensors[self._label_key]
    label = _process_label(label, self._one_hot_label, self._num_classes)

    if self._output_audio:
      audio = decoded_tensors[self._audio_feature]
      audio = t.cast(audio, dtype=self._dtype)
      audio = preprocess_ops_3d.sample_sequence(
          audio, 20, random=False, stride=1)
      audio = tf.ensure_shape(audio, [20, 2048])
      features['audio'] = audio

    return features, label
//  [python]
def _parse_eval_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for evaluation."""
    image = decoded_tensors[self._image_key]
    image = _process_image(
        image=image,
        is_training=False,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
        crop_size=self._crop_size,
        num_crops=self._num_crops,
        zero_centering_image=self._zero_centering_image)
    image = tf.cast(image, dtype=self._dtype)
    features = {'image': image}

    label = decoded_tensors[self._label_key]
    label = _process_label(label, self._one_hot_label, self._num_classes)

    if self._output_audio:
      audio = decoded_tensors[self._audio_feature]
      audio = tf.cast(audio, dtype=self._dtype)
      audio = preprocess_ops_3d.sample_sequence(
          audio, 20, random=False, stride=1)
      audio = tf.ensure_shape(audio, [20, 2048])
      features['audio'] = audio

    return features, label
//  [python]
def get_random_transformation_batch(
        self, batch_size, images=None, **kwargs
    ):
        heights, widths = augmentation_utils.get_images_shape(images)

        tops = tf.where(
            heights < self.height,
            tf.to_float((self.height - heights) / 2),
            tf.zeros_like(heights, dtype=heights.dtype),
        )
        bottoms = tf.where(
            heights < self.height,
            self.height - heights - tops,
            tf.zeros_like(heights, dtype=heights.dtype),
        )
        lefts = tf.where(
            widths < self.width,
            tf.to_float((self.width - widths) / 2),
            tf.zeros_like(widths, dtype=widths.dtype),
        )
        rights = tf.where(
            widths < self.width,
            self.width - widths - lefts,
            tf.zeros_like(widths, dtype=widths.dtype),
        )

        (tops, bottoms, lefts, rights) = augmentation_utils.get_position_params(
            tops, bottoms, lefts, rights, self.position, self._random_generator
        )

        return {
            "pad_tops": tops,
            "pad_bottoms": bottoms,
            "pad_lefts": lefts,
            "pad_rights": rights,
        }
//  [python]
def get_random_transformation_batch(
        self, batch_size, images=None, **kwargs
    ):
        heights, widths = augmentation_utils.get_images_shape(images)

        tops = tf.where(
            heights < self.height,
            tf.cast((self.height - heights) / 2, heights.dtype),
            tf.zeros_like(heights, dtype=heights.dtype),
        )
        bottoms = tf.where(
            heights < self.height,
            self.height - heights - tops,
            tf.zeros_like(heights, dtype=heights.dtype),
        )
        lefts = tf.where(
            widths < self.width,
            tf.cast((self.width - widths) / 2, widths.dtype),
            tf.zeros_like(widths, dtype=widths.dtype),
        )
        rights = tf.where(
            widths < self.width,
            self.width - widths - lefts,
            tf.zeros_like(widths, dtype=widths.dtype),
        )

        (tops, bottoms, lefts, rights) = augmentation_utils.get_position_params(
            tops, bottoms, lefts, rights, self.position, self._random_generator
        )

        return {
            "pad_tops": tops,
            "pad_bottoms": bottoms,
            "pad_lefts": lefts,
            "pad_rights": rights,
        }
//  [python]
def augment_images(self, images, transformations, **kwargs):
        ori_height = tf.shape(images)[augmentation_utils.H_AXIS]
        ori_width = tf.shape(images)[augmentation_utils.W_AXIS]

        pad_top = transformations["pad_tops"][0][0]
        pad_bottom = transformations["pad_bottoms"][0][0]
        pad_left = transformations["pad_lefts"][0][0]
        pad_right = transformations["pad_rights"][0][0]
        paddings = tf.stack(
            (
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
                tf.stack((pad_top, pad_bottom)),
                tf.stack((pad_left, pad_right)),
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
            )
        )
        images = tf.pad(
            images, paddings=paddings, constant_values=self.padding_value
        )

        # center crop
        offset_height = (ori_height + pad_top + pad_bottom - self.height) // 2
        offset_width = (ori_width + pad_left + pad_right - self.width) // 2
        images = tf.image.crop_to_bounding_box(
            images,
            offset_height,
            offset_width,
            self.height,
            self.width,
        )
        images = tf.ensure_shape(
            images, shape=(None, self.height, self.width, None)
        )
        return tf.to_float(images)
//  [python]
def augment_images(self, images, transformations, **kwargs):
        ori_height = tf.shape(images)[augmentation_utils.H_AXIS]
        ori_width = tf.shape(images)[augmentation_utils.W_AXIS]

        pad_top = transformations["pad_tops"][0][0]
        pad_bottom = transformations["pad_bottoms"][0][0]
        pad_left = transformations["pad_lefts"][0][0]
        pad_right = transformations["pad_rights"][0][0]
        paddings = tf.stack(
            (
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
                tf.stack((pad_top, pad_bottom)),
                tf.stack((pad_left, pad_right)),
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
            )
        )
        images = tf.pad(
            images, paddings=paddings, constant_values=self.padding_value
        )

        # center crop
        offset_height = (ori_height + pad_top + pad_bottom - self.height) // 2
        offset_width = (ori_width + pad_left + pad_right - self.width) // 2
        images = tf.image.crop_to_bounding_box(
            images,
            offset_height,
            offset_width,
            self.height,
            self.width,
        )
        images = tf.ensure_shape(
            images, shape=(None, self.height, self.width, None)
        )
        return tf.cast(images, dtype=self.compute_dtype)
//  [python]
def preprocess(obs, config):
    dtype = prec.global_policy().compute_dtype
    obs = obs.copy()
    with tf.device('cpu:0'):
        if 'image' in obs:
            obs['image'] = t.cast(obs['image'], dtype) / 255.0 - 0.5
        if 'lidar' in obs:
            obs['lidar'] = t.cast(obs['lidar'], dtype) / 15.0 - 0.5
        if 'lidar_occupancy' in obs:
            # note: when using `lidar_occupancy` the reconstruction models return a Bernoulli distribution
            # for this reason, we don't center the observation in 0, but let it in [0, 1]
            obs['lidar_occupancy'] = t.cast(obs['lidar_occupancy'], dtype)
        if 'reward' in obs:
            clip_rewards = dict(none=lambda x: x, tanh=tf.tanh,
                                clip=lambda x: tf.clip_by_value(x, config.clip_rewards_min, config.clip_rewards_max))[
                config.clip_rewards]
            obs['reward'] = clip_rewards(obs['reward'])
    return obs
//  [python]
def preprocess(obs, config):
    dtype = prec.global_policy().compute_dtype
    obs = obs.copy()
    with tf.device('cpu:0'):
        if 'image' in obs:
            obs['image'] = tf.cast(obs['image'], dtype) / 255.0 - 0.5
        if 'lidar' in obs:
            obs['lidar'] = tf.cast(obs['lidar'], dtype) / 15.0 - 0.5
        if 'lidar_occupancy' in obs:
            # note: when using `lidar_occupancy` the reconstruction models return a Bernoulli distribution
            # for this reason, we don't center the observation in 0, but let it in [0, 1]
            obs['lidar_occupancy'] = tf.cast(obs['lidar_occupancy'], dtype)
        if 'reward' in obs:
            clip_rewards = dict(none=lambda x: x, tanh=tf.tanh,
                                clip=lambda x: tf.clip_by_value(x, config.clip_rewards_min, config.clip_rewards_max))[
                config.clip_rewards]
            obs['reward'] = clip_rewards(obs['reward'])
    return obs
//  [python]
def augment_segmentation_masks(
        self, segmentation_masks, transformations, **kwargs
    ):
        ori_height = tf.shape(segmentation_masks)[augmentation_utils.H_AXIS]
        ori_width = tf.shape(segmentation_masks)[augmentation_utils.W_AXIS]
        pad_top = transformations["pad_tops"][0][0]
        pad_bottom = transformations["pad_bottoms"][0][0]
        pad_left = transformations["pad_lefts"][0][0]
        pad_right = transformations["pad_rights"][0][0]
        paddings = tf.stack(
            (
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
                tf.stack((pad_top, pad_bottom)),
                tf.stack((pad_left, pad_right)),
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
            )
        )
        segmentation_masks = tf.pad(
            segmentation_masks, paddings=paddings, constant_values=0
        )
        # center crop
        offset_height = (ori_height + pad_top + pad_bottom - self.height) // 2
        offset_width = (ori_width + pad_left + pad_right - self.width) // 2
        segmentation_masks = tf.image.crop_to_bounding_box(
            segmentation_masks,
            offset_height,
            offset_width,
            self.height,
            self.width,
        )
        segmentation_masks = tf.ensure_shape(
            segmentation_masks, shape=(None, self.height, self.width, None)
        )
        return tf.to_float(segmentation_masks)
//  [python]
def augment_segmentation_masks(
        self, segmentation_masks, transformations, **kwargs
    ):
        ori_height = tf.shape(segmentation_masks)[augmentation_utils.H_AXIS]
        ori_width = tf.shape(segmentation_masks)[augmentation_utils.W_AXIS]
        pad_top = transformations["pad_tops"][0][0]
        pad_bottom = transformations["pad_bottoms"][0][0]
        pad_left = transformations["pad_lefts"][0][0]
        pad_right = transformations["pad_rights"][0][0]
        paddings = tf.stack(
            (
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
                tf.stack((pad_top, pad_bottom)),
                tf.stack((pad_left, pad_right)),
                tf.zeros(shape=(2,), dtype=pad_top.dtype),
            )
        )
        segmentation_masks = tf.pad(
            segmentation_masks, paddings=paddings, constant_values=0
        )
        # center crop
        offset_height = (ori_height + pad_top + pad_bottom - self.height) // 2
        offset_width = (ori_width + pad_left + pad_right - self.width) // 2
        segmentation_masks = tf.image.crop_to_bounding_box(
            segmentation_masks,
            offset_height,
            offset_width,
            self.height,
            self.width,
        )
        segmentation_masks = tf.ensure_shape(
            segmentation_masks, shape=(None, self.height, self.width, None)
        )
        return tf.cast(segmentation_masks, dtype=self.compute_dtype)
//  [python]
def augment_bounding_boxes(
        self,
        bounding_boxes,
        transformations,
        images=None,
        raw_images=None,
        **kwargs,
    ):
        if self.bounding_box_format is None:
            raise ValueError(
                "`CenterCrop()` was called with bounding boxes,"
                "but no `bounding_box_format` was specified in the constructor."
                "Please specify a bounding box format in the constructor. i.e."
                "`CenterCrop(..., bounding_box_format='xyxy')`"
            )
        bounding_boxes = bounding_box.to_dense(bounding_boxes)
        bounding_boxes = bounding_box.convert_format(
            bounding_boxes,
            source=self.bounding_box_format,
            target="xyxy",
            images=raw_images,
        )
        original_bounding_boxes = bounding_boxes.copy()

        x1s, y1s, x2s, y2s = tf.split(bounding_boxes["boxes"], 4, axis=-1)
        pad_tops = tf.cast(transformations["pad_tops"], dtype=tf.float32)
        pad_lefts = tf.cast(transformations["pad_lefts"], dtype=tf.float32)
        pad_bottoms = tf.cast(transformations["pad_bottoms"], dtype=tf.float32)
        pad_rights = tf.cast(transformations["pad_rights"], dtype=tf.float32)
        heights, widths = augmentation_utils.get_images_shape(
            raw_images, dtype=tf.float32
        )
        offset_heights = (heights + pad_tops + pad_bottoms - self.height) // 2
        offset_widths = (widths + pad_lefts + pad_rights - self.width) // 2

        x1s += tf.expand_dims(pad_lefts - offset_widths, axis=1)
        y1s += tf.expand_dims(pad_tops - offset_heights, axis=1)
        x2s += tf.expand_dims(pad_lefts - offset_widths, axis=1)
        y2s += tf.expand_dims(pad_tops - offset_heights, axis=1)
        outputs = tf.concat([x1s, y1s, x2s, y2s], axis=-1)

        bounding_boxes = bounding_boxes.copy()
        bounding_boxes["boxes"] = outputs
        bounding_boxes = bounding_box.clip_to_image(
            bounding_boxes,
            bounding_box_format="xyxy",
            images=images,
        )
        bounding_boxes = bounding_box.sanitize_bounding_boxes(
            bounding_boxes,
            min_area_ratio=self.bounding_box_min_area_ratio,
            max_aspect_ratio=self.bounding_box_max_aspect_ratio,
            bounding_box_format="xyxy",
            reference_bounding_boxes=original_bounding_boxes,
            images=images,
            reference_images=raw_images,
        )
        bounding_boxes = bounding_box.convert_format(
            bounding_boxes,
            source="xyxy",
            target=self.bounding_box_format,
            dtype=self.compute_dtype,
            images=images,
        )
        return bounding_boxes
//  [python]
def _valid_range_ids(last_id, capacity):
  """Returns the [min_val, max_val) range of ids.

  Args:
    last_id: A tensor that indicates the last id stored in the replay buffer.
    capacity: The maximum number of elements that the replay buffer can hold.

  Returns:
    A tuple (min_id, max_id) for the range [min_id, max_id) of valid ids.
  """
  min_id_non_full = tf.constant(0, dtype=tf.int64)
  max_id_non_full = tf.maximum(last_id + 1, 0)

  min_id_full = tf.to_int64(last_id + 1 - capacity)
  max_id_full = tf.to_int64(last_id + 1)

  return (
      tf.where(last_id < capacity, min_id_non_full, min_id_full),
      tf.where(last_id < capacity, max_id_non_full, max_id_full),
  )
//  [python]
def _valid_range_ids(last_id, capacity):
  """Returns the [min_val, max_val) range of ids.

  Args:
    last_id: A tensor that indicates the last id stored in the replay buffer.
    capacity: The maximum number of elements that the replay buffer can hold.

  Returns:
    A tuple (min_id, max_id) for the range [min_id, max_id) of valid ids.
  """
  min_id_non_full = tf.constant(0, dtype=tf.int64)
  max_id_non_full = tf.maximum(last_id + 1, 0)

  min_id_full = tf.cast(last_id + 1 - capacity, dtype=tf.int64)
  max_id_full = tf.cast(last_id + 1, dtype=tf.int64)

  return (
      tf.where(last_id < capacity, min_id_non_full, min_id_full),
      tf.where(last_id < capacity, max_id_non_full, max_id_full),
  )
//  [python]
def _sample_episode_ids(
      self, shape, weigh_by_episode_length=False, seed=None
  ):
    """Samples episode ids from the replay buffer."""
    last_id = self._get_last_episode_id()
    assert_nonempty = tf.compat.v1.assert_non_negative(
        last_id,
        message=(
            'EpisodicReplayBuffer is empty. Make sure to add items '
            'before sampling the buffer.'
        ),
    )
    if weigh_by_episode_length:
      # Sample episodes proportional to length.
      with tf.control_dependencies([assert_nonempty]):
        num_episodes = tf.minimum(self._last_episode + 1, self._capacity)
        episode_lengths = self._episode_lengths[:num_episodes]
        logits = tf.math.log(tf.cast(episode_lengths, tf.float32))
        return tf.reshape(
            tf.random.categorical(
                [logits],  # shape is: [1, num_episodes]
                num_samples=tf.reduce_prod(shape),
                seed=seed,
                dtype=tf.int64,
            ),
            shape,
        )
    else:
      min_val, max_val = _valid_range_ids(
          self._get_last_episode_id(), self._capacity
      )
      with tf.control_dependencies([assert_nonempty]):
        return tf.random.uniform(
            shape, minval=min_val, maxval=max_val, dtype=tf.int64, seed=seed
        )
//  [python]
def add_sequence(self, items, episode_id):
    """Adds a sequence of items to the replay buffer for the selected episode.

    Args:
      items: A sequence of items to be added to the buffer. Items will have the
        same structure as the data_spec of this class, but the tensors in items
        will have an outer sequence dimension in addition to the corresponding
        spec in data_spec.
      episode_id: A Tensor containing the current episode_id.

    Returns:
      An updated episode id Tensor.  Accessing this episode id value will,
      as a side effect, start or end the current episode in the buffer.
    """
    episode_id.shape.assert_has_rank(0)
    with tf.device(self._device):
      with tf.name_scope('add_steps'):
        # If users pass in, e.g., a python list [2, 3, 4] of type int32
        # but the data_spec requires an int64, then the user will get a very
        # confusing error much deeper in the TensorList code.  Doing the
        # conversion here either converts when necessary, or raises an error
        # on incompatible types earlier in the run.
        items = tf.nest.map_structure(
            lambda x, spec: tf.convert_to_tensor(value=x, dtype=spec.dtype),
            items,
            self._data_spec,
        )
        item_0 = tf.nest.flatten(items)[0]
        num_steps = tf.cast(
            tf.compat.dimension_value(item_0.shape[0])
            or tf.shape(input=item_0)[0],
            tf.int64,
        )
        # If begin_episode is True, then the increment of the episode_id happens
        # before trying to add anything to the buffer, regardless of whether the
        # item will actually be added.
        begin_episode = self._begin_episode_fn(items)
        end_episode = self._end_episode_fn(items)
        new_episode_id = self._get_episode_id(
            episode_id, begin_episode, end_episode
        )
        episode_location = self._get_episode_id_location(new_episode_id)

        def _add_steps():
          """Add sequence of items to the buffer."""
          inc_episode_length = self._increment_episode_length_locked(
              episode_location, num_steps
          )
          write_data_op = self._data_table.append(episode_location, items)
          with tf.control_dependencies([inc_episode_length, write_data_op]):
            return tf.identity(new_episode_id)

        # Accessing episode_id may modify
        # self._episodes_loc_to_id_map, so ensure it is executed
        # before the tf.equal.
        with tf.control_dependencies([new_episode_id]):
          episode_valid = tf.equal(
              self._episodes_loc_to_id_map[episode_location], new_episode_id
          )

        def _maybe_add_steps():
          return self._add_episode_critical_section.execute(_add_steps)

        return utils.smart_cond(
            episode_valid,
            _maybe_add_steps,
            lambda: tf.identity(new_episode_id),
            name='conditioned_add_steps',
        )
//  [python]
def extract(self, locations, clear_data=False):
    """Extracts Episodes with the given IDs.

    Args:
      locations: A `1-D` Tensor of locations to extract from (note, this is NOT
        the same as an episode ids variable).  It's up to the user to ensure
        that only valid episode locations are requested (i.e., values should be
        between `0` and `self.capacity`). Passing locations that are out of
        bounds will lead to runtime errors.
      clear_data: If `True`, clears the extracted data from this buffer.

    Returns:
      episodes: An Episodes object with an outer dimension of the same size as
        locations.
    """
    locations = tf.cast(locations, dtype=tf.int64, name='locations')
    locations.shape.assert_has_rank(1)

    def _extract_locked():
      """Does the above within the buffer's critical section."""
      episodes = Episodes(
          length=self._episode_lengths.sparse_read(locations),
          completed=self._episode_completed.sparse_read(locations),
          tensor_lists=self._data_table.get_episode_lists(locations),
      )
      if clear_data:
        with tf.control_dependencies(tf.nest.flatten(episodes)):
          clear_rows = self._data_table.clear_rows(locations)
          clear_lengths = tf.compat.v1.scatter_update(
              self._episode_lengths, locations, 0
          )
          clear_completed = tf.compat.v1.scatter_update(
              self._episode_completed, locations, 0
          )
        with tf.control_dependencies(
            [clear_rows, clear_lengths, clear_completed]
        ):
          episodes = tf.nest.map_structure(tf.identity, episodes)
      return episodes

    with tf.device(self._device):
      with tf.name_scope('extract'):
        return self._add_episode_critical_section.execute(_extract_locked)
//  [python]
def test_output_equivalence(self):
        for tokenizer, tf_tokenizer in zip(self.tokenizers, self.tf_tokenizers):
            for test_inputs in (self.test_sentences, self.paired_sentences):
                python_outputs = tokenizer(test_inputs, return_tensors="tf", padding="longest")
                tf_outputs = tf_tokenizer(test_inputs)

                for key in python_outputs.keys():
                    self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape))
                    self.assertTrue(tf.reduce_all(tf.to_int64(python_outputs[key]) == tf_outputs[key]))
//  [python]
def test_output_equivalence(self):
        for tokenizer, tf_tokenizer in zip(self.tokenizers, self.tf_tokenizers):
            for test_inputs in (self.test_sentences, self.paired_sentences):
                python_outputs = tokenizer(test_inputs, return_tensors="tf", padding="longest")
                tf_outputs = tf_tokenizer(test_inputs)

                for key in python_outputs.keys():
                    self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape))
                    self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key], tf.int64) == tf_outputs[key]))
//  [python]
def test_different_pairing_styles(self):
        for tf_tokenizer in self.tf_tokenizers:
            merged_outputs = tf_tokenizer(self.paired_sentences)
            separated_outputs = tf_tokenizer(
                text=[sentence[0] for sentence in self.paired_sentences],
                text_pair=[sentence[1] for sentence in self.paired_sentences],
            )
            for key in merged_outputs.keys():
                self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key], tf.int64) == separated_outputs[key]))
//  [python]
def _gather_all(self):
    """Returns all the items currently in the buffer.

    Returns:
      A tuple containing two entries:
        - All the items currently in the buffer (nested).
        - The items ids.

    Raises:
      ValueError: If the data spec contains lists that must be converted to
        tuples.
    """
    if tf.nest.flatten(self._data_spec) != data_nest.flatten(self._data_spec):
      raise ValueError(
          'Cannot perform gather; data spec contains lists and this conflicts '
          'with gathering operator.  Convert any lists to tuples.  '
          'For example, if your spec looks like [a, b, c], '
          'change it to (a, b, c).  Spec structure is:\n  %s'
          % tf.nest.map_structure(lambda spec: spec.dtype, self._data_spec)
      )

    min_val, max_val = _valid_range_ids(
        self._get_last_episode_id(), self._capacity
    )

    def get_episode_and_id(id_):
      row = self._get_episode_id_location(id_)
      data = self._data_table.get_episode_values(row)
      n = tf.shape(tf.nest.flatten(data)[0])[0]
      id_repeated = tf.fill([n], id_)
      return (tuple(tf.nest.flatten(data)), id_repeated)

    episode_lengths = self._episode_lengths.read_value()
    if self._completed_only:
      episode_lengths *= t.cast(self._episode_completed, dtype=tf.int64)
    total_length = tf.reduce_sum(input_tensor=episode_lengths)

    def via_iterator():
      """If total_length > 0, create a dataset iterator to concat episodes."""
      valid_episodes = tf.range(min_val, max_val)
      ds = tf.data.Dataset.from_tensor_slices(valid_episodes)

      if self._completed_only:
        # Filter out incomplete episodes.
        def check_completed(id_):
          return t.cast(self._episode_completed.sparse_read(id_), tf.bool)

        ds = ds.filter(check_completed)

      def _unflatten(flat_data, id_):
        return tf.nest.pack_sequence_as(self._data_spec, flat_data), id_

      ds = (
          ds.map(get_episode_and_id)
          .unbatch()
          # Batch all the frames in the buffer.  Request a larger amount in
          # case the buffer grows between the construction of total_length and
          # the call to .map().
          .batch(10 + 2 * total_length)
          .map(_unflatten)
      ).batch(1)

      # Use ds.take(1) in case we haven't requested a large enough batch (the
      # replay buffer has grown too quickly), since get_single_element requires
      # that the dataset only has a single entry; and we don't consider this
      # case to be an error.
      return tf.data.experimental.get_single_element(ds.take(1))

    def empty():
      def _empty_from_spec(spec):
        return tf.zeros([0] + spec.shape.as_list(), spec.dtype, name='empty')

      empty_data = tf.nest.map_structure(_empty_from_spec, self._data_spec)
      empty_id = tf.zeros([], dtype=tf.int64)
      return empty_data, empty_id

    return tf.cond(pred=total_length > 0, true_fn=via_iterator, false_fn=empty)
//  [python]
def _gather_all(self):
    """Returns all the items currently in the buffer.

    Returns:
      A tuple containing two entries:
        - All the items currently in the buffer (nested).
        - The items ids.

    Raises:
      ValueError: If the data spec contains lists that must be converted to
        tuples.
    """
    if tf.nest.flatten(self._data_spec) != data_nest.flatten(self._data_spec):
      raise ValueError(
          'Cannot perform gather; data spec contains lists and this conflicts '
          'with gathering operator.  Convert any lists to tuples.  '
          'For example, if your spec looks like [a, b, c], '
          'change it to (a, b, c).  Spec structure is:\n  %s'
          % tf.nest.map_structure(lambda spec: spec.dtype, self._data_spec)
      )

    min_val, max_val = _valid_range_ids(
        self._get_last_episode_id(), self._capacity
    )

    def get_episode_and_id(id_):
      row = self._get_episode_id_location(id_)
      data = self._data_table.get_episode_values(row)
      n = tf.shape(tf.nest.flatten(data)[0])[0]
      id_repeated = tf.fill([n], id_)
      return (tuple(tf.nest.flatten(data)), id_repeated)

    episode_lengths = self._episode_lengths.read_value()
    if self._completed_only:
      episode_lengths *= tf.cast(self._episode_completed, dtype=tf.int64)
    total_length = tf.reduce_sum(input_tensor=episode_lengths)

    def via_iterator():
      """If total_length > 0, create a dataset iterator to concat episodes."""
      valid_episodes = tf.range(min_val, max_val)
      ds = tf.data.Dataset.from_tensor_slices(valid_episodes)

      if self._completed_only:
        # Filter out incomplete episodes.
        def check_completed(id_):
          return tf.cast(self._episode_completed.sparse_read(id_), tf.bool)

        ds = ds.filter(check_completed)

      def _unflatten(flat_data, id_):
        return tf.nest.pack_sequence_as(self._data_spec, flat_data), id_

      ds = (
          ds.map(get_episode_and_id)
          .unbatch()
          # Batch all the frames in the buffer.  Request a larger amount in
          # case the buffer grows between the construction of total_length and
          # the call to .map().
          .batch(10 + 2 * total_length)
          .map(_unflatten)
      ).batch(1)

      # Use ds.take(1) in case we haven't requested a large enough batch (the
      # replay buffer has grown too quickly), since get_single_element requires
      # that the dataset only has a single entry; and we don't consider this
      # case to be an error.
      return tf.data.experimental.get_single_element(ds.take(1))

    def empty():
      def _empty_from_spec(spec):
        return tf.zeros([0] + spec.shape.as_list(), spec.dtype, name='empty')

      empty_data = tf.nest.map_structure(_empty_from_spec, self._data_spec)
      empty_id = tf.zeros([], dtype=tf.int64)
      return empty_data, empty_id

    return tf.cond(pred=total_length > 0, true_fn=via_iterator, false_fn=empty)
//  [python]
def _bf16_to_float32(x: tf.Tensor) -> tf.Tensor:
    # Convert bfloat16 to float32 for numpy compatibility
    return tf.cast(x, tf.float32) if x.dtype == tf.bfloat16 else x
//  [python]
def solarize_add(image, addition=0, threshold=128):
  added_image = t.cast(image, t.int64) + addition
  added_image = t.cast(t.clip_by_value(added_image, 0, 255), t.uint8)
  return t.where(image < threshold, added_image, image)
//  [python]
def solarize_add(image, addition=0, threshold=128):
  # For each pixel in the image less than threshold
  # we add 'addition' amount to it and then clip the
  # pixel value to be between 0 and 255. The value
  # of 'addition' is between -128 and 128.
  added_image = tf.cast(image, tf.int64) + addition
  added_image = tf.cast(tf.clip_by_value(added_image, 0, 255), tf.uint8)
  return tf.where(image < threshold, added_image, image)
//  [python]
def contrast(image, factor):
  """Equivalent of PIL Contrast."""
  degenerate = tf.image.rgb_to_grayscale(image)
  # Cast before calling tf.histogram.
  degenerate = t.cast(degenerate, tf.int32)

  # Compute the grayscale histogram, then compute the mean pixel value,
  # and create a constant image size of that value.  Use that as the
  # blending degenerate target of the original image.
  hist = tf.histogram_fixed_width(degenerate, [0, 255], nbins=256)
  mean = tf.reduce_sum(tf.cast(hist, tf.float32)) / 256.0
  degenerate = tf.ones_like(degenerate, dtype=tf.float32) * mean
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.image.grayscale_to_rgb(t.cast(degenerate, tf.uint8))
  return blend(degenerate, image, factor)
//  [python]
def contrast(image, factor):
  """Equivalent of PIL Contrast."""
  degenerate = tf.image.rgb_to_grayscale(image)
  # Cast before calling tf.histogram.
  degenerate = tf.cast(degenerate, tf.int32)

  # Compute the grayscale histogram, then compute the mean pixel value,
  # and create a constant image size of that value.  Use that as the
  # blending degenerate target of the original image.
  hist = tf.histogram_fixed_width(degenerate, [0, 255], nbins=256)
  mean = tf.reduce_sum(tf.cast(hist, tf.float32)) / 256.0
  degenerate = tf.ones_like(degenerate, dtype=tf.float32) * mean
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.image.grayscale_to_rgb(tf.cast(degenerate, tf.uint8))
  return blend(degenerate, image, factor)
//  [python]
def augment(left_img, right_img):
    active = tf.random_uniform(shape=[4], minval=0, maxval=1, dtype=tf.float32)
    left_img = tf.to_float(left_img)
    right_img = tf.to_float(right_img)

    # random gamma
    # random_gamma = tf.random_uniform(shape=(),minval=0.95,maxval=1.05,dtype=tf.float32)
    # left_img = tf.where(active[0]>0.5,left_img,tf.image.adjust_gamma(left_img,random_gamma))
    # right_img = tf.where(active[0]>0.5,right_img,tf.image.adjust_gamma(right_img,random_gamma))

    # random brightness
    random_delta = tf.random_uniform(shape=(), minval=-0.05, maxval=0.05, dtype=tf.float32)
    left_img = tf.where(active[1] > 0.5, left_img, tf.image.adjust_brightness(left_img, random_delta))
    right_img = tf.where(active[1] > 0.5, right_img, tf.image.adjust_brightness(right_img, random_delta))

    # random contrast
    random_contrast = tf.random_uniform(shape=(), minval=0.8, maxval=1.2, dtype=tf.float32)
    left_img = tf.where(active[2] > 0.5, left_img, tf.image.adjust_contrast(left_img, random_contrast))
    right_img = tf.where(active[2] > 0.5, right_img, tf.image.adjust_contrast(right_img, random_contrast))

    # random hue
    random_hue = tf.random_uniform(shape=(), minval=0.8, maxval=1.2, dtype=tf.float32)
    left_img = tf.where(active[3] > 0.5, left_img,tf.image.adjust_hue(left_img, random_hue))
    right_img = tf.where(active[3] > 0.5, right_img, tf.image.adjust_hue(right_img, random_hue))

    left_img = tf.clip_by_value(left_img,0,255)
    right_img = tf.clip_by_value(right_img,0,255)

    return left_img,right_img
//  [python]
def augment(left_img, right_img):
    active = tf.random_uniform(shape=[4], minval=0, maxval=1, dtype=tf.float32)
    left_img = tf.cast(left_img,tf.float32)
    right_img = tf.cast(right_img,tf.float32)

    # random gamma
    # random_gamma = tf.random_uniform(shape=(),minval=0.95,maxval=1.05,dtype=tf.float32)
    # left_img = tf.where(active[0]>0.5,left_img,tf.image.adjust_gamma(left_img,random_gamma))
    # right_img = tf.where(active[0]>0.5,right_img,tf.image.adjust_gamma(right_img,random_gamma))

    # random brightness
    random_delta = tf.random_uniform(shape=(), minval=-0.05, maxval=0.05, dtype=tf.float32)
    left_img = tf.where(active[1] > 0.5, left_img, tf.image.adjust_brightness(left_img, random_delta))
    right_img = tf.where(active[1] > 0.5, right_img, tf.image.adjust_brightness(right_img, random_delta))

    # random contrast
    random_contrast = tf.random_uniform(shape=(), minval=0.8, maxval=1.2, dtype=tf.float32)
    left_img = tf.where(active[2] > 0.5, left_img, tf.image.adjust_contrast(left_img, random_contrast))
    right_img = tf.where(active[2] > 0.5, right_img, tf.image.adjust_contrast(right_img, random_contrast))

    # random hue
    random_hue = tf.random_uniform(shape=(), minval=0.8, maxval=1.2, dtype=tf.float32)
    left_img = tf.where(active[3] > 0.5, left_img,tf.image.adjust_hue(left_img, random_hue))
    right_img = tf.where(active[3] > 0.5, right_img, tf.image.adjust_hue(right_img, random_hue))

    left_img = tf.clip_by_value(left_img,0,255)
    right_img = tf.clip_by_value(right_img,0,255)

    return left_img,right_img
//  [python]
def _rescale_tf(img,out_shape):
	"""
	Rescale image using bilinear upsampling
	"""
	#print(out_shape)
	def _build_coords(immy,out_shape):
		batch_size = tf.shape(immy)[0]
		in_height = tf.to_float(tf.shape(immy)[1])-1
		in_width = tf.to_float(tf.shape(immy)[2])-1

		out_height = out_shape[0]
		out_width = out_shape[1]

		delta_x = in_width/tf.to_float(out_width-1)
		delta_y = in_height/tf.to_float(out_height-1)

		coord_x = tf.concat([tf.range(in_width-1E-4,delta=delta_x,dtype=tf.float32),[in_width]],axis=0)
		coord_x = tf.expand_dims(coord_x,axis=0)
		coord_x_tile = tf.tile(coord_x,[out_height,1])

		coord_y = tf.concat([tf.range(in_height-1E-4,delta=delta_y,dtype=tf.float32),[in_height]],axis=0)
		coord_y = tf.expand_dims(coord_y,axis=1)
		coord_y_tile = tf.tile(coord_y,[1,out_width])

		coord = tf.stack([coord_x_tile,coord_y_tile],axis=-1)
		#coord = tf.Print(coord,[coord[:,:,0]],summarize=1000)
		coord = tf.expand_dims(coord,axis=0)
		coord = tf.tile(coord,[batch_size,1,1,1])

		return coord
	
	coord = _build_coords(img,out_shape)
	warped = bilinear_sampler(img,coord)
	input_shape = img.get_shape().as_list()
	warped.set_shape([input_shape[0],None,None,input_shape[-1]])
	return warped
//  [python]
def _rescale_tf(img,out_shape):
	"""
	Rescale image using bilinear upsampling
	"""
	#print(out_shape)
	def _build_coords(immy,out_shape):
		batch_size = tf.shape(immy)[0]
		in_height = tf.cast(tf.shape(immy)[1],tf.float32)-1
		in_width = tf.cast(tf.shape(immy)[2],tf.float32)-1

		out_height = out_shape[0]
		out_width = out_shape[1]

		delta_x = in_width/tf.cast(out_width-1,tf.float32)
		delta_y = in_height/tf.cast(out_height-1,tf.float32)

		coord_x = tf.concat([tf.range(in_width-1E-4,delta=delta_x,dtype=tf.float32),[in_width]],axis=0)
		coord_x = tf.expand_dims(coord_x,axis=0)
		coord_x_tile = tf.tile(coord_x,[out_height,1])

		coord_y = tf.concat([tf.range(in_height-1E-4,delta=delta_y,dtype=tf.float32),[in_height]],axis=0)
		coord_y = tf.expand_dims(coord_y,axis=1)
		coord_y_tile = tf.tile(coord_y,[1,out_width])

		coord = tf.stack([coord_x_tile,coord_y_tile],axis=-1)
		#coord = tf.Print(coord,[coord[:,:,0]],summarize=1000)
		coord = tf.expand_dims(coord,axis=0)
		coord = tf.tile(coord,[batch_size,1,1,1])

		return coord
	
	coord = _build_coords(img,out_shape)
	warped = bilinear_sampler(img,coord)
	input_shape = img.get_shape().as_list()
	warped.set_shape([input_shape[0],None,None,input_shape[-1]])
	return warped
//  [python]
def _posterize_level_to_arg(level):
  return (t.cast(level / _MAX_LEVEL * 4, t.uint8),)
//  [python]
def _posterize_level_to_arg(level):
  return (tf.cast(level / _MAX_LEVEL * 4, tf.uint8),)
//  [python]
def _solarize_level_to_arg(level):
  return (t.cast(level / _MAX_LEVEL * 256, t.uint8),)
//  [python]
def _solarize_level_to_arg(level):
  return (tf.cast(level / _MAX_LEVEL * 256, tf.uint8),)
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = t.cast(t.reduce_min(image), t.float32)
    hi = t.cast(t.reduce_max(image), t.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = t.cast(im, t.float32) * scale + offset
      im = t.clip_by_value(im, 0.0, 255.0)
      return t.cast(im, t.uint8)

    result = t.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = t.stack([s1, s2, s3], 2)
  return image
//  [python]
def _solarize_add_level_to_arg(level):
  return (t.cast(level / _MAX_LEVEL * 110, t.int64),)
//  [python]
def _solarize_add_level_to_arg(level):
  return (tf.cast(level / _MAX_LEVEL * 110, tf.int64),)
//  [python]
def sharpness(image, factor):
  """Implements Sharpness function from PIL using TF ops."""
  orig_image = image
  image = t.cast(image, tf.float32)
  # Make image 4D for conv operation.
  image = tf.expand_dims(image, 0)
  # SMOOTH PIL Kernel.
  kernel = tf.constant([[1, 1, 1], [1, 5, 1], [1, 1, 1]],
                       dtype=tf.float32,
                       shape=[3, 3, 1, 1]) / 13.
  # Tile across channel dimension.
  kernel = tf.tile(kernel, [1, 1, 3, 1])
  strides = [1, 1, 1, 1]
  with tf.device('/cpu:0'):
    # Some augmentation that uses depth-wise conv will cause crashing when
    # training on GPU. See (b/156242594) for details.
    degenerate = tf.nn.depthwise_conv2d(image, kernel, strides, padding='VALID')
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.squeeze(t.cast(degenerate, tf.uint8), [0])

  # For the borders of the resulting image, fill in the values of the
  # original image.
  mask = tf.ones_like(degenerate)
  padded_mask = tf.pad(mask, [[1, 1], [1, 1], [0, 0]])
  padded_degenerate = tf.pad(degenerate, [[1, 1], [1, 1], [0, 0]])
  result = tf.where(tf.equal(padded_mask, 1), padded_degenerate, orig_image)

  # Blend the final result.
  return blend(result, orig_image, factor)
//  [python]
def sharpness(image, factor):
  """Implements Sharpness function from PIL using TF ops."""
  orig_image = image
  image = tf.cast(image, tf.float32)
  # Make image 4D for conv operation.
  image = tf.expand_dims(image, 0)
  # SMOOTH PIL Kernel.
  kernel = tf.constant([[1, 1, 1], [1, 5, 1], [1, 1, 1]],
                       dtype=tf.float32,
                       shape=[3, 3, 1, 1]) / 13.
  # Tile across channel dimension.
  kernel = tf.tile(kernel, [1, 1, 3, 1])
  strides = [1, 1, 1, 1]
  with tf.device('/cpu:0'):
    # Some augmentation that uses depth-wise conv will cause crashing when
    # training on GPU. See (b/156242594) for details.
    degenerate = tf.nn.depthwise_conv2d(image, kernel, strides, padding='VALID')
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.squeeze(tf.cast(degenerate, tf.uint8), [0])

  # For the borders of the resulting image, fill in the values of the
  # original image.
  mask = tf.ones_like(degenerate)
  padded_mask = tf.pad(mask, [[1, 1], [1, 1], [0, 0]])
  padded_degenerate = tf.pad(degenerate, [[1, 1], [1, 1], [0, 0]])
  result = tf.where(tf.equal(padded_mask, 1), padded_degenerate, orig_image)

  # Blend the final result.
  return blend(result, orig_image, factor)
//  [python]
def _cutout_arg(level, cutout_size):
  pad_size = t.cast(level / _MAX_LEVEL * cutout_size, t.int32)
  return (2 * pad_size, 2 * pad_size)
//  [python]
def _cutout_arg(level, cutout_size):
  pad_size = tf.cast(level / _MAX_LEVEL * cutout_size, tf.int32)
  return (2 * pad_size, 2 * pad_size)
//  [python]
def equalize(image):
  """Implements Equalize function from PIL using TF ops."""

  def scale_channel(im, c):
    """Scale the data in the channel to implement equalize."""
    im = tf.to_int32(im[:, :, c])
    # Compute the histogram of the image channel.
    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)

    # For the purposes of computing the step, filter out the nonzeros.
    nonzero = tf.where(tf.not_equal(histo, 0))
    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

    def build_lut(histo, step):
      # Compute the cumulative sum, shifting by step // 2
      # and then normalization by step.
      lut = (tf.cumsum(histo) + (step // 2)) // step
      # Shift lut, prepending with 0.
      lut = tf.concat([[0], lut[:-1]], 0)
      # Clip the counts to be in range.  This is done
      # in the C code for image.point.
      return tf.clip_by_value(lut, 0, 255)

    # If step is zero, return the original image.  Otherwise, build
    # lut from the full histogram and step and then index from it.
    result = tf.cond(
        tf.equal(step, 0), lambda: im,
        lambda: tf.gather(build_lut(histo, step), im))

    return tf.to_int8(result)

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image, 0)
  s2 = scale_channel(image, 1)
  s3 = scale_channel(image, 2)
  image = tf.stack([s1, s2, s3], 2)
  return image