jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_symbols_to_logits_fn(self, max_decode_length):
    """Returns a decoding function that calculates logits of the next tokens."""
    timing_signal = self.position_embedding(
        inputs=None, length=max_decode_length + 1)
    timing_signal = tf.cast(timing_signal, dtype=self.compute_dtype)
    decoder_self_attention_mask = tf.linalg.band_part(
        tf.ones([max_decode_length, max_decode_length],
                dtype=self.compute_dtype), -1, 0)
    decoder_self_attention_mask = tf.reshape(
        decoder_self_attention_mask, [1, max_decode_length, max_decode_length])

    def symbols_to_logits_fn(ids, i, cache):
      """Generate logits for next potential IDs.

      Args:
        ids: Current decoded sequences. int tensor with shape `(batch_size *
          beam_size, i + 1)`.
        i: Loop index.
        cache: Dictionary of values storing the encoder output, encoder-decoder
          attention bias, and previous decoder attention values.

      Returns:
        Tuple of
          (logits with shape `(batch_size * beam_size, vocab_size)`,
           updated cache values)
      """
      # Set decoder input to the last generated IDs
      decoder_input = ids[:, -1:]

      # Preprocess decoder input by getting embeddings and adding timing signal.
      decoder_input = self.embedding_lookup(decoder_input)
      decoder_input += timing_signal[i]
      if self._padded_decode:
        # indexing does not work on TPU.
        bias_shape = decoder_self_attention_mask.shape.as_list()
        self_attention_mask = tf.slice(decoder_self_attention_mask, [0, i, 0],
                                       [bias_shape[0], 1, bias_shape[2]])
      else:
        self_attention_mask = decoder_self_attention_mask[:, i:i + 1, :i + 1]
      decoder_shape = tf_utils.get_shape_list(decoder_input, expected_rank=3)
      batch_size = decoder_shape[0]
      decoder_length = decoder_shape[1]

      self_attention_mask = tf.tile(self_attention_mask, [batch_size, 1, 1])
      attention_mask = cache.get("encoder_decoder_attention_mask")
      attention_mask = tf.tile(attention_mask, [1, decoder_length, 1])

      decoder_outputs = self.decoder_layer(
          decoder_input,
          cache.get("encoder_outputs"),
          self_attention_mask=self_attention_mask,
          cross_attention_mask=attention_mask,
          cache=cache,
          decode_loop_step=i if self._padded_decode else None)

      decoder_outputs = tf.cast(decoder_outputs, dtype=self.compute_dtype)
      logits = self._embedding_linear(self.embedding_lookup.embeddings,
                                      decoder_outputs)
      logits = tf.squeeze(logits, axis=[1])
      return logits, cache

    return symbols_to_logits_fn
//  [python]
def data_augment(self, data):
        data_xyz = data[:, :, 0:3]
        batch_size = tf.shape(data_xyz)[0]
        aug_option = np.random.choice([0, 1, 2])
        if aug_option == 0:
            # mirror
            data_xyz = tf.stack([data_xyz[:, :, 0], -data_xyz[:, :, 1], data_xyz[:, :, 2]], 2)
        elif aug_option == 1:
            # rotation
            theta = 2 * 3.14592653 * np.random.rand()
            R = np.array([[np.cos(theta), 0, -np.sin(theta)], [0, 1, 0], [np.sin(theta), 0, np.cos(theta)]])
            R = tf.convert_to_tensor(R, dtype=tf.float32)
            data_xyz = tf.reshape(data_xyz, [-1, 3])
            data_xyz = tf.matmul(data_xyz, R)
            data_xyz = tf.reshape(data_xyz, [-1, self.config.num_points, 3])
        elif aug_option == 2:
            # jitter
            sigma = 0.01
            clip = 0.05
            jittered_point = tf.clip_by_value(sigma * np.random.randn(self.config.num_points, 3), -1 * clip, clip)
            jittered_point = tf.tile(tf.expand_dims(jittered_point, axis=0), [batch_size, 1, 1])
            data_xyz = data_xyz + jittered_point
        if data.get_shape()[-1].value > 3:
            data_f = data[:, :, 3:]
            data_aug = tf.concat([data_xyz, data_f], axis=-1)
        else:
            data_aug = data_xyz
        data_aug_t = tf.transpose(data_aug, [0, 2, 1])
        data_aug_t = tf.reshape(data_aug_t, [-1, data.get_shape()[-1].value, self.config.num_points])
        att_activation = tf.layers.dense(data_aug_t, 1, activation=None, use_bias=False, name='channel_attention')
        att_activation = tf.transpose(att_activation, [0, 2, 1])
        att_scores = tf.nn.softmax(att_activation, axis=-1)
        data_aug = tf.multiply(data_aug, att_scores)
        return data_aug
//  [python]
def data_augment(self, data):
        data_xyz = data[:, :, 0:3]
        batch_size = tf.shape(data_xyz)[0]
        aug_option = np.random.choice([0, 1, 2])
        if aug_option == 0:
            # mirror
            data_xyz = tf.stack([data_xyz[:, :, 0], -data_xyz[:, :, 1], data_xyz[:, :, 2]], 2)
        elif aug_option == 1:
            # rotation
            theta = 2 * 3.14592653 * np.random.rand()
            R = np.array([[np.cos(theta), 0, -np.sin(theta)], [0, 1, 0], [np.sin(theta), 0, np.cos(theta)]])
            R = tf.convert_to_tensor(R, dtype=tf.float32)
            data_xyz = tf.reshape(data_xyz, [-1, 3])
            data_xyz = tf.matmul(data_xyz, R)
            data_xyz = tf.reshape(data_xyz, [-1, self.config.num_points, 3])
        elif aug_option == 2:
            # jitter
            sigma = 0.01
            clip = 0.05
            jittered_point = tf.clip_by_value(sigma * np.random.randn(self.config.num_points, 3), -1 * clip, clip)
            jittered_point = tf.tile(tf.expand_dims(jittered_point, axis=0), [batch_size, 1, 1])
            data_xyz = data_xyz + tf.cast(jittered_point, tf.float32)
        if data.get_shape()[-1].value > 3:
            data_f = data[:, :, 3:]
            data_aug = tf.concat([data_xyz, data_f], axis=-1)
        else:
            data_aug = data_xyz
        data_aug_t = tf.transpose(data_aug, [0, 2, 1])
        data_aug_t = tf.reshape(data_aug_t, [-1, data.get_shape()[-1].value, self.config.num_points])
        att_activation = tf.layers.dense(data_aug_t, 1, activation=None, use_bias=False, name='channel_attention')
        att_activation = tf.transpose(att_activation, [0, 2, 1])
        att_scores = tf.nn.softmax(att_activation, axis=-1)
        data_aug = tf.multiply(data_aug, att_scores)
        return data_aug
//  [python]
def main(unused_argv):
  tf.logging.set_verbosity(tf.logging.INFO)
  tf.logging.info('Prepare to export model to: %s', FLAGS.export_path)

  with tf.Graph().as_default():
    image, image_size, resized_image_size = _create_input_tensors()

    model_options = common.ModelOptions(
        outputs_to_num_classes={common.OUTPUT_TYPE: FLAGS.num_classes},
        crop_size=FLAGS.crop_size,
        atrous_rates=FLAGS.atrous_rates,
        output_stride=FLAGS.output_stride)

    if tuple(FLAGS.inference_scales) == (1.0,):
      tf.logging.info('Exported model performs single-scale inference.')
      predictions = model.predict_labels(
          image,
          model_options=model_options,
          image_pyramid=FLAGS.image_pyramid)
    else:
      tf.logging.info('Exported model performs multi-scale inference.')
      if FLAGS.quantize_delay_step >= 0:
        raise ValueError(
            'Quantize mode is not supported with multi-scale test.')
      predictions = model.predict_labels_multi_scale(
          image,
          model_options=model_options,
          eval_scales=FLAGS.inference_scales,
          add_flipped_images=FLAGS.add_flipped_images)
    raw_predictions = tf.identity(
        tf.cast(predictions[common.OUTPUT_TYPE], tf.float32),
        _RAW_OUTPUT_NAME)
    raw_probabilities = tf.identity(
        predictions[common.OUTPUT_TYPE + model.PROB_SUFFIX],
        _RAW_OUTPUT_PROB_NAME)

    # Crop the valid regions from the predictions.
    semantic_predictions = raw_predictions[
        :, :resized_image_size[0], :resized_image_size[1]]
    semantic_probabilities = raw_probabilities[
        :, :resized_image_size[0], :resized_image_size[1]]

    # Resize back the prediction to the original image size.
    def _resize_label(label, label_size):
      # Expand dimension of label to [1, height, width, 1] for resize operation.
      label = tf.expand_dims(label, 3)
      resized_label = tf.image.resize_images(
          label,
          label_size,
          method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
          align_corners=True)
      return tf.cast(tf.squeeze(resized_label, 3), tf.int32)
    semantic_predictions = _resize_label(semantic_predictions, image_size)
    semantic_predictions = tf.identity(semantic_predictions, name=_OUTPUT_NAME)

    semantic_probabilities = tf.image.resize_bilinear(
        semantic_probabilities, image_size, align_corners=True,
        name=_OUTPUT_PROB_NAME)

    if FLAGS.quantize_delay_step >= 0:
      contrib_quantize.create_eval_graph()

    saver = tf.train.Saver(tf.all_variables())

    dirname = os.path.dirname(FLAGS.export_path)
    tf.gfile.MakeDirs(dirname)
    graph_def = tf.get_default_graph().as_graph_def(add_shapes=True)
    freeze_graph.freeze_graph_with_def_protos(
        graph_def,
        saver.as_saver_def(),
        FLAGS.checkpoint_path,
        _OUTPUT_NAME + ',' + _OUTPUT_PROB_NAME,
        restore_op_name=None,
        filename_tensor_name=None,
        output_graph=FLAGS.export_path,
        clear_devices=True,
        initializer_nodes=None)

    if FLAGS.save_inference_graph:
      tf.train.write_graph(graph_def, dirname, 'inference_graph.pbtxt')
//  [python]
def call(self, inp):
        if self.div_val == 1:
            raise NotImplementedError  # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
        else:
            inp_flat = tf.reshape(inp, (-1,))
            emb_flat = tf.zeros([shape_list(inp_flat)[0], self.d_proj])
            for i in range(len(self.cutoffs)):
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]

                mask_i = (inp_flat >= l_idx) & (inp_flat < r_idx)

                inp_i = tf.boolean_mask(inp_flat, mask_i) - l_idx
                emb_i = self.emb_layers[i](inp_i)
                emb_i = tf.einsum("id,de->ie", emb_i, self.emb_projs[i])

                mask_idx = tf.where(mask_i)
                scatter = tf.scatter_nd(mask_idx, emb_i, shape_list(emb_flat))
                emb_flat = tf.cast(emb_flat, dtype=scatter.dtype)
                emb_flat += scatter

            embed_shape = shape_list(inp) + [self.d_proj]
            embed = tf.reshape(emb_flat, embed_shape)

        embed *= self.emb_scale

        return embed
//  [python]
def call(self, inputs: tf.Tensor, step_type=None, network_state=()):
    del step_type
    inputs = tf.cast(inputs, tf.float32)
    for layer in self._dummy_layers:
      inputs = layer(inputs)
    return inputs, network_state
//  [python]
def _update_mems(self, hids, mems, mlen, qlen):
        # does not deal with None
        if mems is None:
            return None

        # mems is not None
        assert len(hids) == len(mems), "len(hids) != len(mems)"

        # There are `mlen + qlen` steps that can be cached into mems
        new_mems = []
        end_idx = mlen + tf.math.maximum(0, qlen)
        beg_idx = tf.math.maximum(0, end_idx - tf.convert_to_tensor(self.mem_len))
        for i in range(len(hids)):
            mems[i] = t.cast(mems[i], dtype=hids[i].dtype)
            cat = tf.concat([mems[i], hids[i]], axis=0)
            tf.stop_gradient(cat)
            new_mems.append(cat[beg_idx:end_idx])

        return new_mems
//  [python]
def _update_mems(self, hids, mems, mlen, qlen):
        # does not deal with None
        if mems is None:
            return None

        # mems is not None
        assert len(hids) == len(mems), "len(hids) != len(mems)"

        # There are `mlen + qlen` steps that can be cached into mems
        new_mems = []
        end_idx = mlen + tf.math.maximum(0, qlen)
        beg_idx = tf.math.maximum(0, end_idx - tf.convert_to_tensor(self.mem_len))
        for i in range(len(hids)):
            mems[i] = tf.cast(mems[i], dtype=hids[i].dtype)
            cat = tf.concat([mems[i], hids[i]], axis=0)
            tf.stop_gradient(cat)
            new_mems.append(cat[beg_idx:end_idx])

        return new_mems
//  [python]
def pair_masks(mask_x, mask_y):
  """Combines a pair of 2D masks into a single 3D mask.

  Args:
    mask_x: A tf.Tensor<float>[batch, len_x] with binary entries.
    mask_y: A tf.Tensor<float>[batch, len_y] with binary entries.

  Returns:
    A tf.Tensor<float>[batch, len_x, len_y] with binary entries, defined as
      out[n][i][j] := mask_x[n][i] * mask_y[n][j].
  """
  mask1, mask2 = t.cast(mask_x, t.float32), t.cast(mask_y, t.float32)
  return t.cast(t.einsum('ij,ik->ijk', mask1, mask2), t.bool)
//  [python]
def pair_masks(mask_x, mask_y):
  """Combines a pair of 2D masks into a single 3D mask.

  Args:
    mask_x: A tf.Tensor<float>[batch, len_x] with binary entries.
    mask_y: A tf.Tensor<float>[batch, len_y] with binary entries.

  Returns:
    A tf.Tensor<float>[batch, len_x, len_y] with binary entries, defined as
      out[n][i][j] := mask_x[n][i] * mask_y[n][j].
  """
  mask1, mask2 = tf.cast(mask_x, tf.float32), tf.cast(mask_y, tf.float32)
  return tf.cast(tf.einsum('ij,ik->ijk', mask1, mask2), tf.bool)
//  [python]
def get_nearest_centers(inputs, centers):
  """Returns the nearest center to each example in input.

  Args:
    inputs: A tf.float32 tensor of size [N, D] where N is the number of data
      points and D is their dimension.
    centers: A tf.float32 tensor of size [K, D] where K is the number of
      centers.

  Returns:
    A [N] tf.int32 tensor of nearest center indices.
  """
  input_distances = inputs_distances_to_centers(inputs, centers)
  return tf.to_int32(tf.math.argmin(input_distances, axis=1))
//  [python]
def get_nearest_centers(inputs, centers):
  """Returns the nearest center to each example in input.

  Args:
    inputs: A tf.float32 tensor of size [N, D] where N is the number of data
      points and D is their dimension.
    centers: A tf.float32 tensor of size [K, D] where K is the number of
      centers.

  Returns:
    A [N] tf.int32 tensor of nearest center indices.
  """
  input_distances = inputs_distances_to_centers(inputs, centers)
  return tf.cast(tf.math.argmin(input_distances, axis=1), dtype=tf.int32)
//  [python]
def _create_arm_policy_and_observations(
      self,
  ) -> Tuple[
      greedy_multi_objective_policy.GreedyMultiObjectiveNeuralPolicy,
      Dict[Text, tf.Tensor],
  ]:
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(2, 3, 4)
    time_step_spec = ts.time_step_spec(obs_spec)
    action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 3)
    objective_networks = [
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
        for _ in range(3)
    ]
    policy = greedy_multi_objective_policy.GreedyMultiObjectiveNeuralPolicy(
        time_step_spec,
        action_spec,
        self._scalarizer,
        objective_networks,
        accepts_per_arm_features=True,
        emit_policy_info=(
            policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
            policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,
        ),
    )
    action_feature = t.cast(
        tf.reshape(tf.random.shuffle(tf.range(24)), shape=[2, 4, 3]),
        dtype=tf.float32,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [2, 1]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: action_feature,
    }
    return policy, observations
//  [python]
def _create_arm_policy_and_observations(
      self,
  ) -> Tuple[
      greedy_multi_objective_policy.GreedyMultiObjectiveNeuralPolicy,
      Dict[Text, tf.Tensor],
  ]:
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(2, 3, 4)
    time_step_spec = ts.time_step_spec(obs_spec)
    action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 3)
    objective_networks = [
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
        for _ in range(3)
    ]
    policy = greedy_multi_objective_policy.GreedyMultiObjectiveNeuralPolicy(
        time_step_spec,
        action_spec,
        self._scalarizer,
        objective_networks,
        accepts_per_arm_features=True,
        emit_policy_info=(
            policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
            policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,
        ),
    )
    action_feature = tf.cast(
        tf.reshape(tf.random.shuffle(tf.range(24)), shape=[2, 4, 3]),
        dtype=tf.float32,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [2, 1]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: action_feature,
    }
    return policy, observations
//  [python]
def map_labels_to_0_to_n(labels):
  """Maps the values in labels to numbers between 0 and N.

  Finds the unique values in labels, and maps them to new values.
  The goal is to make sure that the output labels has values starting at 0
  and ending at N. It can be that some pixels have value (-1) which will remain
  as they are.

  Args:
    labels: Input tensor of tf.int32.

  Returns:
    Mapped labels, same size as input labels with values between 0 and N.
  """
  labels_shape = labels.get_shape()
  tf_labels_shape = tf.shape(labels)
  labels = tf.reshape(labels, [-1])
  mask_ignore = tf.less(labels, 0)
  mask_no_ignore = 1 - tf.cast(mask_ignore, dtype=tf.int32)
  labels = labels * tf.cast(mask_no_ignore, dtype=tf.int32)
  max_label = tf.reduce_max(labels)
  labels += tf.cast(mask_ignore, dtype=tf.int32) * max_label
  label_map = tf.zeros([max_label + 1], tf.int32)
  unique_labels, _ = tf.unique(labels)
  label_map = tf.reduce_sum(
      tf.one_hot(unique_labels, depth=max_label+1, dtype=tf.float32), axis=0)
  label_map = tf.cast((tf.cumsum(label_map) - 1.0) * label_map, dtype=tf.int32)
  labels = tf.map_fn(lambda l: label_map[l], labels, back_prop=False)
  labels = (labels * tf.cast(mask_no_ignore, dtype=tf.int32)) - tf.cast(mask_ignore, dtype=tf.int32)
  labels = tf.reshape(labels, tf_labels_shape)
  labels.set_shape(labels_shape)
  return labels
//  [python]
def map_labels_to_0_to_n(labels):
  """Maps the values in labels to numbers between 0 and N.

  Finds the unique values in labels, and maps them to new values.
  The goal is to make sure that the output labels has values starting at 0
  and ending at N. It can be that some pixels have value (-1) which will remain
  as they are.

  Args:
    labels: Input tensor of tf.int32.

  Returns:
    Mapped labels, same size as input labels with values between 0 and N.
  """
  labels_shape = labels.get_shape()
  tf_labels_shape = tf.shape(labels)
  labels = tf.reshape(labels, [-1])
  mask_ignore = tf.cast(tf.less(labels, 0), dtype=tf.int32)
  mask_no_ignore = 1 - mask_ignore
  labels *= mask_no_ignore
  max_label = tf.reduce_max(labels)
  labels += mask_ignore * max_label
  label_map = tf.zeros([max_label + 1], tf.int32)
  unique_labels, _ = tf.unique(labels)
  label_map = tf.reduce_sum(
      tf.one_hot(unique_labels, depth=max_label+1, dtype=tf.float32), axis=0)
  label_map = tf.cast((tf.cumsum(label_map) - 1.0) * label_map, dtype=tf.int32)
  labels = tf.map_fn(lambda l: label_map[l], labels, back_prop=False)
  labels = (labels * mask_no_ignore) - mask_ignore
  labels = tf.reshape(labels, tf_labels_shape)
  labels.set_shape(labels_shape)
  return labels
//  [python]
def randomly_select_n_points_per_segment(labels, num_points,
                                         include_ignore_label=False):
  """Randomly selects a point for each segment.

  Args:
    labels: A tf.int32 tensor of [N] containing labels for pixels. labels start
            at 0.
    num_points: Number of samples.
    include_ignore_label: If this flag is true, the function will not ignore
                          the pixels with ignore label (-1). It increases the
                          labels by 1 and will sample from ignore label too.

  Returns:
    indices: Indices of randomly picked centers. A tf.int32 tensor of size
             [num_points, num_classes].
  """
  if include_ignore_label:
    labels -= tf.reduce_min(labels)
  i = tf.constant(0, dtype=tf.int32)
  num_classes = tf.reduce_max(labels) + 1
  indices = tf.zeros(tf.stack([num_points, num_classes]), dtype=tf.int32)
  indices_shape = indices.get_shape()

  def body(i, indices):
    indices_local, _ = randomly_select_one_point_per_segment(labels)
    pad0 = tf.stack([i, num_points-i-1])
    pad1 = tf.constant([0, 0], dtype=tf.int32)
    padding = tf.stack([pad0, pad1])
    indices_local = tf.pad(tf.expand_dims(indices_local, 0), padding)
    indices += t.cast(indices_local, dtype=tf.int32)
    indices.set_shape(indices_shape)
    i += 1
    return (i, indices)

  (i, indices) = tf.while_loop(lambda i, indices: tf.less(i, num_points),
                               body, [i, indices])
  return indices
//  [python]
def randomly_select_n_points_per_segment(labels, num_points,
                                         include_ignore_label=False):
  """Randomly selects a point for each segment.

  Args:
    labels: A tf.int32 tensor of [N] containing labels for pixels. labels start
            at 0.
    num_points: Number of samples.
    include_ignore_label: If this flag is true, the function will not ignore
                          the pixels with ignore label (-1). It increases the
                          labels by 1 and will sample from ignore label too.

  Returns:
    indices: Indices of randomly picked centers. A tf.int32 tensor of size
             [num_points, num_classes].
  """
  if include_ignore_label:
    labels -= tf.reduce_min(labels)
  i = tf.constant(0, dtype=tf.int32)
  num_classes = tf.reduce_max(labels) + 1
  indices = tf.zeros(tf.stack([num_points, num_classes]), dtype=tf.int32)
  indices_shape = indices.get_shape()

  def body(i, indices):
    indices_local, _ = randomly_select_one_point_per_segment(labels)
    pad0 = tf.stack([i, num_points-i-1])
    pad1 = tf.constant([0, 0], dtype=tf.int32)
    padding = tf.stack([pad0, pad1])
    indices_local = tf.pad(tf.expand_dims(indices_local, 0), padding)
    indices += tf.cast(indices_local, dtype=tf.int32)
    indices.set_shape(indices_shape)
    i += 1
    return (i, indices)

  (i, indices) = tf.while_loop(lambda i, indices: tf.less(i, num_points),
                               body, [i, indices])
  return indices
//  [python]
def call(self, w, r, attn_mask, mems, head_mask, output_attentions, training=False):
        qlen, rlen, bsz = shape_list(w)[0], shape_list(r)[0], shape_list(w)[1]

        if mems is not None:
            mems = t.cast(mems, dtype=w.dtype)
            cat = tf.concat([mems, w], 0)
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(cat))
            else:
                w_heads = self.qkv_net(cat)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, axis=-1)
            w_head_q = w_head_q[-qlen:]
        else:
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(w))
            else:
                w_heads = self.qkv_net(w)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, axis=-1)

        klen = shape_list(w_head_k)[0]

        w_head_q = tf.reshape(w_head_q, (qlen, bsz, self.n_head, self.d_head))  # qlen x bsz x n_head x d_head
        w_head_k = tf.reshape(w_head_k, (klen, bsz, self.n_head, self.d_head))  # qlen x bsz x n_head x d_head
        w_head_v = tf.reshape(w_head_v, (klen, bsz, self.n_head, self.d_head))  # qlen x bsz x n_head x d_head

        r_head_k = tf.reshape(r_head_k, (rlen, self.n_head, self.d_head))  # qlen x n_head x d_head

        # compute attention score
        rw_head_q = w_head_q + self.r_w_bias  # qlen x bsz x n_head x d_head
        AC = tf.einsum("ibnd,jbnd->ijbn", rw_head_q, w_head_k)  # qlen x klen x bsz x n_head

        rr_head_q = w_head_q + self.r_r_bias
        BD = tf.einsum("ibnd,jnd->ijbn", rr_head_q, r_head_k)  # qlen x klen x bsz x n_head
        BD = self._rel_shift(BD)

        # [qlen x klen x bsz x n_head]
        attn_score = AC + BD
        attn_score = attn_score * self.scale

        # compute attention probability
        if attn_mask is not None:
            attn_mask_t = attn_mask[:, :, None, None]
            attn_mask_t = tf.cast(attn_mask_t, dtype=attn_score.dtype)
            attn_score = attn_score * (1.0 - attn_mask_t) - 1e30 * attn_mask_t

        # [qlen x klen x bsz x n_head]
        attn_prob = stable_softmax(attn_score, axis=1)
        attn_prob = self.dropatt(attn_prob, training=training)

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

        # compute attention vector
        attn_vec = tf.einsum("ijbn,jbnd->ibnd", attn_prob, w_head_v)

        # [qlen x bsz x n_head x d_head]
        attn_vec_sizes = shape_list(attn_vec)
        attn_vec = tf.reshape(attn_vec, (attn_vec_sizes[0], attn_vec_sizes[1], self.n_head * self.d_head))

        # linear projection
        attn_out = self.o_net(attn_vec)
        attn_out = self.drop(attn_out, training=training)

        if self.pre_lnorm:
            # residual connection
            outputs = [w + attn_out]
        else:
            # residual connection + layer normalization
            outputs = [self.layer_norm(w + attn_out)]

        if output_attentions:
            outputs.append(attn_prob)

        return outputs
//  [python]
def call(self, w, r, attn_mask, mems, head_mask, output_attentions, training=False):
        qlen, rlen, bsz = shape_list(w)[0], shape_list(r)[0], shape_list(w)[1]

        if mems is not None:
            mems = tf.cast(mems, dtype=w.dtype)
            cat = tf.concat([mems, w], 0)
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(cat))
            else:
                w_heads = self.qkv_net(cat)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, axis=-1)
            w_head_q = w_head_q[-qlen:]
        else:
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(w))
            else:
                w_heads = self.qkv_net(w)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, axis=-1)

        klen = shape_list(w_head_k)[0]

        w_head_q = tf.reshape(w_head_q, (qlen, bsz, self.n_head, self.d_head))  # qlen x bsz x n_head x d_head
        w_head_k = tf.reshape(w_head_k, (klen, bsz, self.n_head, self.d_head))  # qlen x bsz x n_head x d_head
        w_head_v = tf.reshape(w_head_v, (klen, bsz, self.n_head, self.d_head))  # qlen x bsz x n_head x d_head

        r_head_k = tf.reshape(r_head_k, (rlen, self.n_head, self.d_head))  # qlen x n_head x d_head

        # compute attention score
        rw_head_q = w_head_q + self.r_w_bias  # qlen x bsz x n_head x d_head
        AC = tf.einsum("ibnd,jbnd->ijbn", rw_head_q, w_head_k)  # qlen x klen x bsz x n_head

        rr_head_q = w_head_q + self.r_r_bias
        BD = tf.einsum("ibnd,jnd->ijbn", rr_head_q, r_head_k)  # qlen x klen x bsz x n_head
        BD = self._rel_shift(BD)

        # [qlen x klen x bsz x n_head]
        attn_score = AC + BD
        attn_score = attn_score * self.scale

        # compute attention probability
        if attn_mask is not None:
            attn_mask_t = attn_mask[:, :, None, None]
            attn_mask_t = tf.cast(attn_mask_t, dtype=attn_score.dtype)
            attn_score = attn_score * (1.0 - attn_mask_t) - 1e30 * attn_mask_t

        # [qlen x klen x bsz x n_head]
        attn_prob = stable_softmax(attn_score, axis=1)
        attn_prob = self.dropatt(attn_prob, training=training)

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

        # compute attention vector
        attn_vec = tf.einsum("ijbn,jbnd->ibnd", attn_prob, w_head_v)

        # [qlen x bsz x n_head x d_head]
        attn_vec_sizes = shape_list(attn_vec)
        attn_vec = tf.reshape(attn_vec, (attn_vec_sizes[0], attn_vec_sizes[1], self.n_head * self.d_head))

        # linear projection
        attn_out = self.o_net(attn_vec)
        attn_out = self.drop(attn_out, training=training)

        if self.pre_lnorm:
            # residual connection
            outputs = [w + attn_out]
        else:
            # residual connection + layer normalization
            outputs = [self.layer_norm(w + attn_out)]

        if output_attentions:
            outputs.append(attn_prob)

        return outputs
//  [python]
def randomly_select_one_point_per_segment(labels, include_ignore_label=False):
  """Randomly selects a point for each segment.

  Args:
    labels: A tf.int32 tensor of [N] containing labels for pixels. labels start
            at 0. It can be that some pixels have label (-1) which means this
            function will ignore those pixels.
    include_ignore_label: If this flag is true, the function will not ignore
                          the pixels with ignore label (-1). It increases the
                          labels by 1 and will sample from ignore label too.

  Returns:
    indices: Indices of randomly picked centers. A tensor of type tf.int32.
    labels_one_hot: One hot of the labels. A tensor of type tf.float32.
  """
  if include_ignore_label:
    labels -= tf.reduce_min(labels)
  num_pixels = tf.size(labels)
  num_classes = tf.reduce_max(labels) + 1
  labels_one_hot = tf.one_hot(labels, num_classes, dtype=tf.float32)
  labels_size = tf.reduce_sum(labels_one_hot, 0)
  labels_cumsum = tf.cumsum(labels_one_hot, axis=0)
  labels_cumsum = t.cast(labels_cumsum * labels_one_hot, dtype=tf.int32)
  rand_inds = tf.random.uniform(
      tf.shape(labels_size), minval=0, maxval=1.0, dtype=tf.float32)
  rand_inds = tf.maximum(
      1, t.cast(rand_inds * (labels_size + 1.0), dtype=tf.int32))
  rand_inds = tf.minimum(rand_inds, t.cast(labels_size, dtype=tf.int32))
  rand_inds = tf.expand_dims(rand_inds, 0)
  rand_inds = tf.tile(rand_inds, tf.stack([num_pixels, 1]))
  cond = tf.equal(labels_cumsum, rand_inds)
  indices = tf.where(tf.transpose(cond))
  _, indices = tf.split(value=indices, num_or_size_splits=2, axis=1)
  indices = tf.reshape(indices, [-1])
  return indices, labels_one_hot
//  [python]
def randomly_select_one_point_per_segment(labels, include_ignore_label=False):
  """Randomly selects a point for each segment.

  Args:
    labels: A tf.int32 tensor of [N] containing labels for pixels. labels start
            at 0. It can be that some pixels have label (-1) which means this
            function will ignore those pixels.
    include_ignore_label: If this flag is true, the function will not ignore
                          the pixels with ignore label (-1). It increases the
                          labels by 1 and will sample from ignore label too.

  Returns:
    indices: Indices of randomly picked centers. A tensor of type tf.int32.
    labels_one_hot: One hot of the labels. A tensor of type tf.float32.
  """
  if include_ignore_label:
    labels -= tf.reduce_min(labels)
  num_pixels = tf.size(labels)
  num_classes = tf.reduce_max(labels) + 1
  labels_one_hot = tf.one_hot(labels, num_classes, dtype=tf.float32)
  labels_size = tf.reduce_sum(labels_one_hot, 0)
  labels_cumsum = tf.cumsum(labels_one_hot, axis=0)
  labels_cumsum = tf.cast(labels_cumsum * labels_one_hot, dtype=tf.int32)
  rand_inds = tf.random.uniform(
      tf.shape(labels_size), minval=0, maxval=1.0, dtype=tf.float32)
  rand_inds = tf.maximum(
      1, tf.cast(rand_inds * (labels_size + 1.0), dtype=tf.int32))
  rand_inds = tf.minimum(rand_inds, tf.cast(labels_size, dtype=tf.int32))
  rand_inds = tf.expand_dims(rand_inds, 0)
  rand_inds = tf.tile(rand_inds, tf.stack([num_pixels, 1]))
  cond = tf.equal(labels_cumsum, rand_inds)
  indices = tf.where(tf.transpose(cond))
  _, indices = tf.split(value=indices, num_or_size_splits=2, axis=1)
  indices = tf.reshape(indices, [-1])
  return indices, labels_one_hot
//  [python]
def points_mask_pairwise_iou(masks1, masks2):
  """Intersection over union between pairwise corresponding masks.

  Args:
    masks1: A tensor of size [num_masks, num_points] with values ranging
      between 0 and 1.
    masks2: A tensor of size [num_masks, num_points] with values ranging
      between 0 and 1.

  Returns:
    A tensor of size [num_masks].
  """
  masks1 = tf.to_float(masks1)
  masks2 = tf.to_float(masks2)
  intersections = tf.math.reduce_sum(masks1 * masks2, axis=1)
  mask_sizes1 = tf.math.reduce_sum(masks1, axis=1)
  mask_sizes2 = tf.math.reduce_sum(masks2, axis=1)
  unions = tf.maximum(mask_sizes1 + mask_sizes2 - intersections, 1.0)
  return intersections / unions
//  [python]
def points_mask_pairwise_iou(masks1, masks2):
  """Intersection over union between pairwise corresponding masks.

  Args:
    masks1: A tensor of size [num_masks, num_points] with values ranging
      between 0 and 1.
    masks2: A tensor of size [num_masks, num_points] with values ranging
      between 0 and 1.

  Returns:
    A tensor of size [num_masks].
  """
  masks1 = tf.cast(masks1, dtype=tf.float32)
  masks2 = tf.cast(masks2, dtype=tf.float32)
  intersections = tf.math.reduce_sum(masks1 * masks2, axis=1)
  mask_sizes1 = tf.math.reduce_sum(masks1, axis=1)
  mask_sizes2 = tf.math.reduce_sum(masks2, axis=1)
  unions = tf.maximum(mask_sizes1 + mask_sizes2 - intersections, 1.0)
  return intersections / unions
//  [python]
def points_mask_iou(masks1, masks2):
  """Intersection over union between point masks.

  Args:
    masks1: A tensor of size [num_masks1, num_points] with values ranging
      between 0 and 1.
    masks2: A tensor of size [num_masks2, num_points] with values ranging
      between 0 and 1.

  Returns:
    A tensor of size [num_masks1, num_masks2].
  """
  num_masks1 = tf.shape(masks1)[0]
  num_masks2 = tf.shape(masks2)[0]
  masks1 = tf.to_float(masks1)
  masks2 = tf.to_float(masks2)
  intersections = tf.linalg.matmul(masks1, masks2, transpose_b=True)
  mask_sizes1 = tf.math.reduce_sum(masks1, axis=[1])
  mask_sizes1 = tf.tile(
      tf.expand_dims(mask_sizes1, axis=1), tf.stack([1, num_masks2]))
  mask_sizes2 = tf.math.reduce_sum(masks2, axis=[1])
  mask_sizes2 = tf.tile(
      tf.expand_dims(mask_sizes2, axis=0), tf.stack([num_masks1, 1]))
  unions = tf.maximum(mask_sizes1 + mask_sizes2 - intersections, 1.0)
  return intersections / unions
//  [python]
def points_mask_iou(masks1, masks2):
  """Intersection over union between point masks.

  Args:
    masks1: A tensor of size [num_masks1, num_points] with values ranging
      between 0 and 1.
    masks2: A tensor of size [num_masks2, num_points] with values ranging
      between 0 and 1.

  Returns:
    A tensor of size [num_masks1, num_masks2].
  """
  num_masks1 = tf.shape(masks1)[0]
  num_masks2 = tf.shape(masks2)[0]
  masks1 = tf.cast(masks1, dtype=tf.float32)
  masks2 = tf.cast(masks2, dtype=tf.float32)
  intersections = tf.linalg.matmul(masks1, masks2, transpose_b=True)
  mask_sizes1 = tf.math.reduce_sum(masks1, axis=[1])
  mask_sizes1 = tf.tile(
      tf.expand_dims(mask_sizes1, axis=1), tf.stack([1, num_masks2]))
  mask_sizes2 = tf.math.reduce_sum(masks2, axis=[1])
  mask_sizes2 = tf.tile(
      tf.expand_dims(mask_sizes2, axis=0), tf.stack([num_masks1, 1]))
  unions = tf.maximum(mask_sizes1 + mask_sizes2 - intersections, 1.0)
  return intersections / unions
//  [python]
def call(self, inputs, step_type=None, network_state=()):
    del step_type
    inputs = t.cast(inputs, tf.float32)
    for layer in self._dummy_layers:
      inputs = layer(inputs)
    return inputs, network_state
//  [python]
def call(self, inputs, step_type=None, network_state=()):
    del step_type
    inputs = tf.cast(inputs, tf.float32)
    for layer in self._dummy_layers:
      inputs = layer(inputs)
    return inputs, network_state
//  [python]
def dummy_inputs(self):
        pad_token = 1
        input_ids = t.convert_to_tensor(DUMMY_INPUTS, dtype=t.int32)
        decoder_input_ids = t.convert_to_tensor(DUMMY_INPUTS, dtype=t.int32)
        dummy_inputs = {
            "decoder_input_ids": decoder_input_ids,
            "attention_mask": tf.math.not_equal(input_ids, pad_token),
            "input_ids": input_ids,
        }
        return dummy_inputs
//  [python]
def dummy_inputs(self):
        pad_token = 1
        input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32)
        decoder_input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32)
        dummy_inputs = {
            "decoder_input_ids": decoder_input_ids,
            "attention_mask": tf.math.not_equal(input_ids, pad_token),
            "input_ids": input_ids,
        }
        return dummy_inputs
//  [python]
def sample_indices(dim_x, size=128, sort=False):
  dim_x = t.cast(dim_x, tf.int32)
  indices = tf.range(0, dim_x, dtype=tf.int32)
  indices = tf.random.shuffle(indices)[:size]
  if sort:
    indices = tf.sort(indices)
  return indices
//  [python]
def sample_indices(dim_x, size=128, sort=False):
  dim_x = tf.cast(dim_x, tf.int32)
  indices = tf.range(0, dim_x, dtype=tf.int32)
  indices = tf.random.shuffle(indices)[:size]
  if sort:
    indices = tf.sort(indices)
  return indices
//  [python]
def _calculate_action_cost_matrix(ac1, ac2):
  diff = tf.expand_dims(ac1, axis=1) - tf.expand_dims(ac2, axis=0)
  return tf.reduce_mean(tf.abs(diff), axis=-1)
//  [python]
def _calculate_action_cost_matrix(ac1, ac2):
  diff = tf.expand_dims(ac1, axis=1) - tf.expand_dims(ac2, axis=0)
  return tf.cast(tf.reduce_mean(tf.abs(diff), axis=-1), dtype=tf.float32)
//  [python]
def contrastive_loss(similarity_matrix,
                     metric_values,
                     temperature,
                     coupling_temperature=1.0,
                     use_coupling_weights=True):
  """Contrative Loss with soft coupling."""
  assert temperature > 0.
  metric_shape = tf.shape(metric_values)
  similarity_matrix /= temperature
  neg_logits1 = similarity_matrix

  col_indices = tf.argmin(metric_values, axis=1)
  pos_indices1 = tf.stack(
      (tf.range(metric_shape[0]), col_indices), axis=1)
  pos_logits1 = tf.gather_nd(similarity_matrix, pos_indices1)

  if use_coupling_weights:
    metric_values /= coupling_temperature
    coupling = tf.exp(-metric_values)
    pos_weights1 = -tf.gather_nd(metric_values, pos_indices1)
    pos_logits1 += pos_weights1
    negative_weights = tf.math.log((1.0 - coupling) + EPS)
    neg_logits1 += tf.tensor_scatter_nd_update(negative_weights, pos_indices1,
                                               pos_weights1)
  neg_logits1 = tf.math.reduce_logsumexp(neg_logits1, axis=1)
  return tf.reduce_mean(neg_logits1 - pos_logits1)
//  [python]
def contrastive_loss(similarity_matrix,
                     metric_values,
                     temperature,
                     coupling_temperature=1.0,
                     use_coupling_weights=True):
  """Contrative Loss with soft coupling."""
  assert temperature > 0.
  metric_shape = tf.shape(metric_values)
  similarity_matrix /= temperature
  neg_logits1 = similarity_matrix

  col_indices = tf.cast(tf.argmin(metric_values, axis=1), dtype=tf.int32)
  pos_indices1 = tf.stack(
      (tf.range(metric_shape[0], dtype=tf.int32), col_indices), axis=1)
  pos_logits1 = tf.gather_nd(similarity_matrix, pos_indices1)

  if use_coupling_weights:
    metric_values /= coupling_temperature
    coupling = tf.exp(-metric_values)
    pos_weights1 = -tf.gather_nd(metric_values, pos_indices1)
    pos_logits1 += pos_weights1
    negative_weights = tf.math.log((1.0 - coupling) + EPS)
    neg_logits1 += tf.tensor_scatter_nd_update(negative_weights, pos_indices1,
                                               pos_weights1)
  neg_logits1 = tf.math.reduce_logsumexp(neg_logits1, axis=1)
  return tf.reduce_mean(neg_logits1 - pos_logits1)
//  [python]
def sample_random_sequence(batch_video_matrix, num_frames, num_samples):
  """Samples a random sequence of frames of size num_samples.

  Args:
    batch_video_matrix: tensor of shape [batch_size x max_frames x feature_size]
    num_frames: tensor of shape [batch_size x 1]
    num_samples: a scalar indicating the number of samples

  Returns:
    reshaped batch_video_matrix in [batch_size x 'num_samples' x feature_size]
  """

  batch_size = tf.shape(batch_video_matrix)[0]
  frame_index_offset = tf.tile(
      tf.expand_dims(tf.range(num_samples), 0), [batch_size, 1])
  max_start_frame_index = tf.maximum(num_frames - num_samples, 0)
  start_frame_index = tf.to_int32(
      tf.multiply(
          tf.random.uniform([batch_size, 1]),
          tf.to_float(max_start_frame_index + 1)))
  frame_index = tf.minimum(start_frame_index + frame_index_offset,
                           tf.to_int32(num_frames - 1))
  batch_index = tf.tile(
      tf.expand_dims(tf.range(batch_size), 1), [1, num_samples])
  index = tf.stack([batch_index, frame_index], 2)
  return tf.gather_nd(batch_video_matrix, index)
//  [python]
def sample_random_sequence(batch_video_matrix, num_frames, num_samples):
  """Samples a random sequence of frames of size num_samples.

  Args:
    batch_video_matrix: tensor of shape [batch_size x max_frames x feature_size]
    num_frames: tensor of shape [batch_size x 1]
    num_samples: a scalar indicating the number of samples

  Returns:
    reshaped batch_video_matrix in [batch_size x 'num_samples' x feature_size]
  """

  batch_size = tf.shape(batch_video_matrix)[0]
  frame_index_offset = tf.tile(
      tf.expand_dims(tf.range(num_samples), 0), [batch_size, 1])
  max_start_frame_index = tf.maximum(num_frames - num_samples, 0)
  start_frame_index = tf.cast(
      tf.multiply(
          tf.random.uniform([batch_size, 1]),
          tf.cast(max_start_frame_index + 1, tf.float32)), tf.int32)
  frame_index = tf.minimum(start_frame_index + frame_index_offset,
                           tf.cast(num_frames - 1, tf.int32))
  batch_index = tf.tile(
      tf.expand_dims(tf.range(batch_size), 1), [1, num_samples])
  index = tf.stack([batch_index, frame_index], 2)
  return tf.gather_nd(batch_video_matrix, index)
//  [python]
def fit_embedding(self, states, actions,
                    next_states, next_actions, rewards,
                    discounts, level_ids):
    """Fit representation (pixel encoder).

    Args:
      states: tf.tensor
      actions: tf.tensor
      next_states: tf.tensor
      next_actions: tf.tensor
      rewards: tf.tensor
      discounts: tf.tensor
      level_ids: tf.tensor, contains level id labels

    Returns:
      embedding_dict: dict

    """
    del next_actions, discounts, rewards, next_states
    ssl_variables = (self.critic.trainable_variables +
                     self.embedding.trainable_variables)

    with tf.GradientTape(watch_accessed_variables=True) as tape:
      uniques, _, counts = tf.unique_with_counts(level_ids)
      uniques = tf.to_int32(uniques)
      level_order = tf.argsort(counts, direction='DESCENDING')
      # Take two most frequent levels in the batch
      lvls = tf.gather(uniques, level_order)[:2]
      lvl_1, lvl_2 = lvls[0], lvls[1]
      min_count = tf.reduce_min(tf.gather(counts, level_order)[:2])

      idx_1 = tf.math.equal(level_ids, lvl_1)
      idx_2 = tf.math.equal(level_ids, lvl_2)

      act1 = tf.one_hot(actions[idx_1][:min_count], 15)
      act2 = tf.one_hot(actions[idx_2][:min_count], 15)

      representation = self.embedding(self.critic.encoder(states[0]))
      representation_1 = representation[idx_1][:min_count]
      representation_2 = representation[idx_2][:min_count]

      metric_vals = compute_metric(act1, act2, self.discount)

      embedding_loss = representation_alignment_loss(
          representation_1,
          representation_2,
          metric_vals,
          use_coupling_weights=False,
          temperature=self.temperature,
          return_representation=False)

    br_grads = tape.gradient(embedding_loss, ssl_variables)
    self.br_optimizer.apply_gradients(zip(br_grads, ssl_variables))

    gn = tf.reduce_mean([tf.linalg.norm(v) for v in br_grads if v is not None])

    return {
        'embedding_loss': embedding_loss,
        'embedding_grad_norm': gn
    }
//  [python]
def fit_embedding(self, states, actions,
                    next_states, next_actions, rewards,
                    discounts, level_ids):
    """Fit representation (pixel encoder).

    Args:
      states: tf.tensor
      actions: tf.tensor
      next_states: tf.tensor
      next_actions: tf.tensor
      rewards: tf.tensor
      discounts: tf.tensor
      level_ids: tf.tensor, contains level id labels

    Returns:
      embedding_dict: dict

    """
    del next_actions, discounts, rewards, next_states
    ssl_variables = (self.critic.trainable_variables +
                     self.embedding.trainable_variables)

    with tf.GradientTape(watch_accessed_variables=True) as tape:
      uniques, _, counts = tf.unique_with_counts(level_ids)
      uniques = tf.cast(uniques, dtype=tf.int32)
      level_order = tf.argsort(counts, direction='DESCENDING')
      # Take two most frequent levels in the batch
      lvls = tf.gather(uniques, level_order)[:2]
      lvl_1, lvl_2 = lvls[0], lvls[1]
      min_count = tf.reduce_min(tf.gather(counts, level_order)[:2])

      idx_1 = tf.math.equal(level_ids, lvl_1)
      idx_2 = tf.math.equal(level_ids, lvl_2)

      act1 = tf.one_hot(actions[idx_1][:min_count], 15)
      act2 = tf.one_hot(actions[idx_2][:min_count], 15)

      representation = self.embedding(self.critic.encoder(states[0]))
      representation_1 = representation[idx_1][:min_count]
      representation_2 = representation[idx_2][:min_count]

      metric_vals = compute_metric(act1, act2, self.discount)

      embedding_loss = representation_alignment_loss(
          representation_1,
          representation_2,
          metric_vals,
          use_coupling_weights=False,
          temperature=self.temperature,
          return_representation=False)

    br_grads = tape.gradient(embedding_loss, ssl_variables)
    self.br_optimizer.apply_gradients(zip(br_grads, ssl_variables))

    gn = tf.reduce_mean([tf.linalg.norm(v) for v in br_grads if v is not None])

    return {
        'embedding_loss': embedding_loss,
        'embedding_grad_norm': gn
    }
//  [python]
def sample_random_frames(batch_video_matrix, num_frames, num_samples):
  """Samples a random set of frames of size num_samples.

  Args:
    batch_video_matrix: tensor of shape [batch_size x max_frames x feature_size]
    num_frames: tensor of shape [batch_size x 1]
    num_samples (int): a scalar indicating the number of samples

  Returns:
    reshaped batch_video_matrix in [batch_size x 'num_samples' x feature_size]
  """
  batch_size = tf.shape(batch_video_matrix)[0]
  frame_index = tf.to_int32(
      tf.multiply(
          tf.random.uniform([batch_size, num_samples]),
          tf.tile(num_frames, [1, num_samples])))
  batch_index = tf.tile(
      tf.expand_dims(tf.range(batch_size), 1), [1, num_samples])
  index = tf.stack([batch_index, frame_index], 2)
  return tf.gather_nd(batch_video_matrix, index)
//  [python]
def sample_random_frames(batch_video_matrix, num_frames, num_samples):
  """Samples a random set of frames of size num_samples.

  Args:
    batch_video_matrix: tensor of shape [batch_size x max_frames x feature_size]
    num_frames: tensor of shape [batch_size x 1]
    num_samples (int): a scalar indicating the number of samples

  Returns:
    reshaped batch_video_matrix in [batch_size x 'num_samples' x feature_size]
  """
  batch_size = tf.shape(batch_video_matrix)[0]
  frame_index = tf.cast(
      tf.multiply(
          tf.random.uniform([batch_size, num_samples]),
          tf.tile(num_frames, [1, num_samples])), tf.int32)
  batch_index = tf.tile(
      tf.expand_dims(tf.range(batch_size), 1), [1, num_samples])
  index = tf.stack([batch_index, frame_index], 2)
  return tf.gather_nd(batch_video_matrix, index)