jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def preprocess_fn(*features, mean, std, augment=False):
  """Preprocess datasets."""
  features = features[0]
  image = features['image']
  label = t.cast(features['label'], tf.int32)
  image = t.cast(image, tf.float32) / 255.0
  if augment:
    image = tf.image.resize_with_crop_or_pad(image, IMAGE_SIZE + 4,
                                             IMAGE_SIZE + 4)
    image = tf.image.random_crop(image,
                                 [image.shape[0], IMAGE_SIZE, IMAGE_SIZE, 3])
    image = tf.image.random_flip_left_right(image)
    # image = tf.image.per_image_standardization(image)
    image = (image - mean) / std
  else:
    image = tf.image.resize_with_crop_or_pad(image, IMAGE_SIZE, IMAGE_SIZE)
    # image = tf.image.per_image_standardization(image)
    image = (image - mean) / std
  return dict(image=image, label=label)
//  [python]
def preprocess_fn(*features, mean, std, augment=False):
  """Preprocess datasets."""
  features = features[0]
  image = features['image']
  label = tf.cast(features['label'], tf.int32)
  image = tf.cast(image, tf.float32) / 255.0
  if augment:
    image = tf.image.resize_with_crop_or_pad(image, IMAGE_SIZE + 4,
                                             IMAGE_SIZE + 4)
    image = tf.image.random_crop(image,
                                 [image.shape[0], IMAGE_SIZE, IMAGE_SIZE, 3])
    image = tf.image.random_flip_left_right(image)
    # image = tf.image.per_image_standardization(image)
    image = (image - mean) / std
  else:
    image = tf.image.resize_with_crop_or_pad(image, IMAGE_SIZE, IMAGE_SIZE)
    # image = tf.image.per_image_standardization(image)
    image = (image - mean) / std
  return dict(image=image, label=label)
//  [python]
def _preproc_test_samples(self, inputs):
        inputs = self._normalize_inputs(inputs)
        return tf.cast(inputs, self.precision)
//  [python]
def padded_one_hot_encoding(indices, depth, left_pad):
  """Returns a zero padded one-hot tensor.

  This function converts a sparse representation of indices (e.g., [4]) to a
  zero padded one-hot representation (e.g., [0, 0, 0, 0, 1] with depth = 4 and
  left_pad = 1). If `indices` is empty, the result will simply be a tensor of
  shape (0, depth + left_pad). If depth = 0, then this function just returns
  `None`.

  Args:
    indices: an integer tensor of shape [num_indices].
    depth: depth for the one-hot tensor (integer).
    left_pad: number of zeros to left pad the one-hot tensor with (integer).

  Returns:
    padded_onehot: a tensor with shape (num_indices, depth + left_pad). Returns
      `None` if the depth is zero.

  Raises:
    ValueError: if `indices` does not have rank 1 or if `left_pad` or `depth are
      either negative or non-integers.

  TODO: add runtime checks for depth and indices.
  """
  if depth < 0 or not isinstance(depth, (int, long)):
    raise ValueError('`depth` must be a non-negative integer.')
  if left_pad < 0 or not isinstance(left_pad, (int, long)):
    raise ValueError('`left_pad` must be a non-negative integer.')
  if depth == 0:
    return None
  if len(indices.get_shape().as_list()) != 1:
    raise ValueError('`indices` must have rank 1')

  def one_hot_and_pad():
    one_hot = tf.to_float(tf.one_hot(tf.to_int64(indices), depth,
                                 on_value=1, off_value=0))
    return tf.pad(one_hot, [[0, 0], [left_pad, 0]], mode='CONSTANT')
  result = tf.cond(tf.greater(tf.size(indices), 0), one_hot_and_pad,
                   lambda: tf.zeros((depth + left_pad, 0)))
  return tf.reshape(result, [-1, depth + left_pad])
//  [python]
def padded_one_hot_encoding(indices, depth, left_pad):
  """Returns a zero padded one-hot tensor.

  This function converts a sparse representation of indices (e.g., [4]) to a
  zero padded one-hot representation (e.g., [0, 0, 0, 0, 1] with depth = 4 and
  left_pad = 1). If `indices` is empty, the result will simply be a tensor of
  shape (0, depth + left_pad). If depth = 0, then this function just returns
  `None`.

  Args:
    indices: an integer tensor of shape [num_indices].
    depth: depth for the one-hot tensor (integer).
    left_pad: number of zeros to left pad the one-hot tensor with (integer).

  Returns:
    padded_onehot: a tensor with shape (num_indices, depth + left_pad). Returns
      `None` if the depth is zero.

  Raises:
    ValueError: if `indices` does not have rank 1 or if `left_pad` or `depth are
      either negative or non-integers.

  TODO: add runtime checks for depth and indices.
  """
  if depth < 0 or not isinstance(depth, (int, long)):
    raise ValueError('`depth` must be a non-negative integer.')
  if left_pad < 0 or not isinstance(left_pad, (int, long)):
    raise ValueError('`left_pad` must be a non-negative integer.')
  if depth == 0:
    return None
  if len(indices.get_shape().as_list()) != 1:
    raise ValueError('`indices` must have rank 1')

  def one_hot_and_pad():
    one_hot = tf.cast(tf.one_hot(tf.cast(indices, tf.int64), depth,
                                 on_value=1, off_value=0), tf.float32)
    return tf.pad(one_hot, [[0, 0], [left_pad, 0]], mode='CONSTANT')
  result = tf.cond(tf.greater(tf.size(indices), 0), one_hot_and_pad,
                   lambda: tf.zeros((depth + left_pad, 0)))
  return tf.reshape(result, [-1, depth + left_pad])
//  [python]
def _preproc_eval_samples(self, inputs, labels):
        """Preprocess samples and perform random augmentations"""
        inputs = self._normalize_inputs(inputs)
        labels = self._normalize_labels(labels)

        # Bring back labels to network's output size and remove interpolation artifacts
        labels = tf.image.resize_with_crop_or_pad(labels, target_width=388, target_height=388)
        cond = tf.less(labels, 0.5 * tf.ones(tf.shape(input=labels)))
        labels = tf.where(cond, tf.zeros(tf.shape(input=labels)), tf.ones(tf.shape(input=labels)))

        return tf.cast(inputs, self.precision), labels
//  [python]
def add_loss_op(self, q, target_q):
        """
        Sets the loss of a batch, self.loss is a scalar

        Args:
            q: (tf tensor) shape = (batch_size, num_actions)
            target_q: (tf tensor) shape = (batch_size, num_actions)
        """
        # you may need this variable
        num_actions = self.env.action_space.n

        ##############################################################
        """
        TODO: 
            The loss for an example is defined as:
                Q_samp(s) = r if done
                          = r + gamma * max_a' Q_target(s', a')
                loss = (Q_samp(s) - Q(s, a))^2 
        HINT: 
            - Config variables are accessible through self.config
            - You can access placeholders like self.a (for actions)
                self.r (rewards) or self.done_mask for instance
            - You may find the following functions useful
                - tf.cast
                - tf.reduce_max
                - tf.reduce_sum
                - tf.one_hot
                - tf.squared_difference
                - tf.reduce_mean
        """
        ##############################################################
        ##################### YOUR CODE HERE - 4-5 lines #############
        
        best_target_q = tf.reduce_max(target_q, axis=1)
        Q_samp = self.r + (1.0 - tf.to_float(self.done_mask)) * self.config.gamma * best_target_q
        Q_sa = tf.reduce_sum(q * tf.one_hot(self.a, self.env.action_space.n), axis=1)
        self.loss = tf.reduce_mean(tf.squared_difference(Q_samp, Q_sa))
//  [python]
def add_loss_op(self, q, target_q):
        """
        Sets the loss of a batch, self.loss is a scalar

        Args:
            q: (tf tensor) shape = (batch_size, num_actions)
            target_q: (tf tensor) shape = (batch_size, num_actions)
        """
        # you may need this variable
        num_actions = self.env.action_space.n

        ##############################################################
        """
        TODO: 
            The loss for an example is defined as:
                Q_samp(s) = r if done
                          = r + gamma * max_a' Q_target(s', a')
                loss = (Q_samp(s) - Q(s, a))^2 
        HINT: 
            - Config variables are accessible through self.config
            - You can access placeholders like self.a (for actions)
                self.r (rewards) or self.done_mask for instance
            - You may find the following functions useful
                - tf.cast
                - tf.reduce_max
                - tf.reduce_sum
                - tf.one_hot
                - tf.squared_difference
                - tf.reduce_mean
        """
        ##############################################################
        ##################### YOUR CODE HERE - 4-5 lines #############
        
        best_target_q = tf.reduce_max(target_q, axis=1)
        Q_samp = self.r + (1.0 - tf.cast(self.done_mask, tf.float32)) * self.config.gamma * best_target_q
        Q_sa = tf.reduce_sum(q * tf.one_hot(self.a, self.env.action_space.n), axis=1)
        self.loss = tf.reduce_mean(tf.squared_difference(Q_samp, Q_sa))
//  [python]
def prepare_led_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
):
    if attention_mask is None:
        attention_mask = tf.to_int32(tf.math.not_equal(input_ids, config.pad_token_id))
    if decoder_attention_mask is None:
        decoder_attention_mask = tf.concat(
            [
                tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int32),
                tf.to_int32(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id)),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = tf.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "decoder_input_ids": decoder_input_ids,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
    }
//  [python]
def prepare_led_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
):
    if attention_mask is None:
        attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8)
    if decoder_attention_mask is None:
        decoder_attention_mask = tf.concat(
            [
                tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8),
                tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = tf.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "decoder_input_ids": decoder_input_ids,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
    }
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFLEDModel(config=config).get_decoder()
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = tf.to_int32(ids_tensor((self.batch_size, 3), 2))

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFLEDModel(config=config).get_decoder()
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8)

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def train_step(self, transitions):
    states, actions, rewards, _, mask = transitions
    states = tf.where(states < 0, 0, states)
    actions = tf.where(actions < 0, 0, actions)
    rewards = tf.where(rewards < 0., 0., rewards)
    states = tf.reshape(states, [-1])
    actions = tf.reshape(actions, [-1])
    rewards = tf.reshape(rewards, [-1])
    rewards = tf.cast(rewards, tf.int32)
    this_rtg = rewards

    with tf.GradientTape() as tape:
      this_rtg = tf.cast(this_rtg, actions.dtype)
      indices = tf.stack([states, actions, this_rtg], axis=-1)
      z_logits = tf.gather_nd(self.posterior_net, indices)
      z_probs = tf.nn.softmax(z_logits)
      z_preds = tf.squeeze(tf.random.categorical(z_logits, 1), -1)
      z_preds = tf.one_hot(z_preds, self.num_latents)
      # Straight through gradient
      z_preds = z_preds + z_probs - tf.stop_gradient(z_probs)

      value_logits = tf.reduce_sum(
          z_preds[Ellipsis, None] * tf.gather(self.value_net, states), axis=1)
      value_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          this_rtg, value_logits)
      action_logits = tf.reduce_sum(
          z_preds[Ellipsis, None] * tf.gather(self.policy, states), axis=1)
      action_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          actions, action_logits)

      pos_indices = tf.stack([states, actions, this_rtg], axis=-1)
      pos_embeds = tf.gather_nd(self.energy_net, pos_indices)
      pos_loss = tf.math.reduce_sum(pos_embeds * z_preds, axis=-1)
      neg_indices = tf.stack([states, actions], axis=-1)
      neg_embeds = tf.gather_nd(self.energy_net, neg_indices)
      neg_loss = tfp.math.reduce_weighted_logsumexp(
          tf.reduce_sum(neg_embeds * z_preds[:, None, :], axis=-1),
          w=self.rtg_prior,
          axis=-1)

      energy_loss = self.energy_weight * (-pos_loss + neg_loss)
      loss = (
          tf.reduce_mean(value_loss) + tf.reduce_mean(action_loss) +
          tf.reduce_mean(energy_loss))

    grads = tape.gradient(loss, self.variables)
    self.optimizer.apply_gradients(zip(grads, self.variables))
    return {
        'loss': loss,
        'value_loss': tf.reduce_mean(value_loss),
        'action_loss': tf.reduce_mean(action_loss),
        'pos_loss': tf.reduce_mean(pos_loss),
        'neg_loss': tf.reduce_mean(neg_loss),
        'energy_loss': tf.reduce_mean(energy_loss),
    }
//  [python]
def filter_episodes_rnn(traj):
  """Map TFRecord windows (of adjacent TimeSteps) so only first episode valid.

  Args:
    traj: Trajectory.

  Returns:
    Trajectory containing filtered sample with only one episode.
  """
  step_types = traj.step_type
  seq_len = t.cast(t.shape(step_types)[0], t.int32)

  # Find the first "LAST" frame.  Set everything after it to be invalid.
  last_frames = t.where(step_types == StepType.LAST)
  if t.shape(last_frames)[0] == 0:
    discount = t.ones_like(traj.discount)
  else:
    first_last = t.cast(last_frames[0], t.int32) + 1
    if first_last == seq_len:
      discount = t.ones_like(traj.discount)
    else:
      valid = t.ones(first_last)
      not_valid = t.zeros(seq_len - first_last)
      discount = t.concat([valid, not_valid], axis=0)

  return traj._replace(discount=discount)
//  [python]
def filter_episodes_rnn(traj):
  """Map TFRecord windows (of adjacent TimeSteps) so only first episode valid.

  Args:
    traj: Trajectory.

  Returns:
    Trajectory containing filtered sample with only one episode.
  """
  step_types = traj.step_type
  seq_len = tf.cast(tf.shape(step_types)[0], tf.int32)

  # Find the first "LAST" frame.  Set everything after it to be invalid.
  last_frames = tf.where(step_types == StepType.LAST)
  if tf.shape(last_frames)[0] == 0:
    discount = tf.ones_like(traj.discount)
  else:
    first_last = tf.cast(last_frames[0], tf.int32) + 1
    if first_last == seq_len:
      discount = tf.ones_like(traj.discount)
    else:
      valid = tf.ones(first_last)
      not_valid = tf.zeros(seq_len - first_last)
      discount = tf.concat([valid, not_valid], axis=0)

  return traj._replace(discount=discount)
//  [python]
def filter_episodes(traj):
  """Map TFRecord windows (of adjacent TimeSteps) to single episodes.

  Outputs the last episode within a sample window. It does this by using
  the step_type tensor to break up sequences into single episode sections.
  For example, if step_type is: [FIRST, MID, LAST, FIRST, MID, MID], we
  will return a sample, whos tensor indices are sampled as:
  [3, 3, 3, 3, 4, 5]. So that the 3rd index frame is replicated 3 times to
  across the beginning of the tensor.

  Args:
    traj: Trajectory.

  Returns:
    Trajectory containing filtered sample with only one episode.
  """
  step_types = traj.step_type
  seq_len = t.cast(t.shape(step_types)[0], t.int32)

  # Find the last start frame in the window. e.g. if we have step types
  # [FIRST, MID, LAST, FIRST, MID, MID], we want index 3.
  first_frames = t.where(step_types == StepType.FIRST)

  if t.shape(first_frames)[0] == 0:
    # No first frame, return sequence as is.
    inds = t.range(0, seq_len)
  else:
    ind_start = t.cast(first_frames[-1, 0], t.int32)
    if ind_start == 0:
      # Last episode starts on the first frame, return as is.
      inds = t.range(0, seq_len)
    else:
      # Otherwise, resample so that the last episode's first frame is
      # replicated to the beginning of the sample. In the example above we want:
      # [3, 3, 3, 3, 4, 5].
      inds_start = t.tile(ind_start[None], ind_start[None])
      inds_end = t.range(ind_start, seq_len)
      inds = t.concat([inds_start, inds_end], axis=0)

  def _resample(arr):
    if isinstance(arr, t.Tensor):
      return t.gather(arr, inds)
    else:
      return arr  # empty or None

  observation = t.nest.map_structure(_resample, traj.observation)

  return Trajectory(
      step_type=_resample(traj.step_type),
      action=_resample(traj.action),
      policy_info=_resample(traj.policy_info),
      next_step_type=_resample(traj.next_step_type),
      reward=_resample(traj.reward),
      discount=_resample(traj.discount),
      observation=observation)
//  [python]
def filter_episodes(traj):
  """Map TFRecord windows (of adjacent TimeSteps) to single episodes.

  Outputs the last episode within a sample window. It does this by using
  the step_type tensor to break up sequences into single episode sections.
  For example, if step_type is: [FIRST, MID, LAST, FIRST, MID, MID], we
  will return a sample, whos tensor indices are sampled as:
  [3, 3, 3, 3, 4, 5]. So that the 3rd index frame is replicated 3 times to
  across the beginning of the tensor.

  Args:
    traj: Trajectory.

  Returns:
    Trajectory containing filtered sample with only one episode.
  """
  step_types = traj.step_type
  seq_len = tf.cast(tf.shape(step_types)[0], tf.int32)

  # Find the last start frame in the window. e.g. if we have step types
  # [FIRST, MID, LAST, FIRST, MID, MID], we want index 3.
  first_frames = tf.where(step_types == StepType.FIRST)

  if tf.shape(first_frames)[0] == 0:
    # No first frame, return sequence as is.
    inds = tf.range(0, seq_len)
  else:
    ind_start = tf.cast(first_frames[-1, 0], tf.int32)
    if ind_start == 0:
      # Last episode starts on the first frame, return as is.
      inds = tf.range(0, seq_len)
    else:
      # Otherwise, resample so that the last episode's first frame is
      # replicated to the beginning of the sample. In the example above we want:
      # [3, 3, 3, 3, 4, 5].
      inds_start = tf.tile(ind_start[None], ind_start[None])
      inds_end = tf.range(ind_start, seq_len)
      inds = tf.concat([inds_start, inds_end], axis=0)

  def _resample(arr):
    if isinstance(arr, tf.Tensor):
      return tf.gather(arr, inds)
    else:
      return arr  # empty or None

  observation = tf.nest.map_structure(_resample, traj.observation)

  return Trajectory(
      step_type=_resample(traj.step_type),
      action=_resample(traj.action),
      policy_info=_resample(traj.policy_info),
      next_step_type=_resample(traj.next_step_type),
      reward=_resample(traj.reward),
      discount=_resample(traj.discount),
      observation=observation)
//  [python]
def _spread_out_weights(self, non_match_dot):
        non_match_dot = tf.reduce_max(tf.reshape(
            non_match_dot, shape=[-1, self.n_negatives]), axis=1)

        m1 = tf.reduce_sum(non_match_dot)
        m2 = tf.reduce_sum(tf.square(non_match_dot))
        n_samples = tf.to_float(tf.shape(self.anchor_ids)[0])
        return tf.maximum(
            (tf.square(m1) + tf.maximum(0.0, m2 - n_samples /
                                        self.embedding_dim)) *
            self.spreadout_weight, 0.)
//  [python]
def _spread_out_weights(self, non_match_dot):
        non_match_dot = tf.reduce_max(tf.reshape(
            non_match_dot, shape=[-1, self.n_negatives]), axis=1)

        m1 = tf.reduce_sum(non_match_dot)
        m2 = tf.reduce_sum(tf.square(non_match_dot))
        n_samples = tf.cast(tf.shape(self.anchor_ids)[0], tf.float32)
        return tf.maximum(
            (tf.square(m1) + tf.maximum(0.0, m2 - n_samples /
                                        self.embedding_dim)) *
            self.spreadout_weight, 0.)
//  [python]
def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int):
    """
    Shift input ids one token to the right, and wrap the last non pad token (the <LID> token) Note that MBart does not
    have a single `decoder_start_token_id` in contrast to other Bart-like models.
    """
    if pad_token_id is None:
        raise ValueError("self.model.config.pad_token_id has to be defined.")
    # replace possible -100 values in labels by `pad_token_id`
    input_ids = tf.where(
        input_ids == -100, tf.fill(shape_list(input_ids), t.cast(pad_token_id, input_ids.dtype)), input_ids
    )
    language_id_index = (
        tf.reduce_sum(tf.cast(tf.math.not_equal(input_ids, pad_token_id), dtype=input_ids.dtype), axis=-1) - 1
    )
    language_id_index = tf.stack(
        [tf.range(shape_list(input_ids)[0], dtype=input_ids.dtype), language_id_index], axis=-1
    )
    languages_ids = tf.gather_nd(input_ids, language_id_index)

    shifted_input_ids = tf.concat([tf.expand_dims(languages_ids, axis=-1), input_ids[:, :-1]], axis=-1)

    return shifted_input_ids
//  [python]
def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int):
    """
    Shift input ids one token to the right, and wrap the last non pad token (the <LID> token) Note that MBart does not
    have a single `decoder_start_token_id` in contrast to other Bart-like models.
    """
    if pad_token_id is None:
        raise ValueError("self.model.config.pad_token_id has to be defined.")
    # replace possible -100 values in labels by `pad_token_id`
    input_ids = tf.where(
        input_ids == -100, tf.fill(shape_list(input_ids), tf.cast(pad_token_id, input_ids.dtype)), input_ids
    )
    language_id_index = (
        tf.reduce_sum(tf.cast(tf.math.not_equal(input_ids, pad_token_id), dtype=input_ids.dtype), axis=-1) - 1
    )
    language_id_index = tf.stack(
        [tf.range(shape_list(input_ids)[0], dtype=input_ids.dtype), language_id_index], axis=-1
    )
    languages_ids = tf.gather_nd(input_ids, language_id_index)

    shifted_input_ids = tf.concat([tf.expand_dims(languages_ids, axis=-1), input_ids[:, :-1]], axis=-1)

    return shifted_input_ids
//  [python]
def test_mask_invalid_locations(self):
        hidden_states = self._get_hidden_states()
        batch_size = 1
        seq_length = 8
        hidden_size = 4
        hidden_states = tf.reshape(hidden_states, (batch_size, seq_length, hidden_size))
        hidden_states = TFLongformerSelfAttention._chunk(hidden_states, window_overlap=2)

        hid_states_1 = TFLongformerSelfAttention._mask_invalid_locations(hidden_states, 1)
        hid_states_2 = TFLongformerSelfAttention._mask_invalid_locations(hidden_states, 2)
        hid_states_3 = TFLongformerSelfAttention._mask_invalid_locations(hidden_states[:, :, :, :3], 2)
        hid_states_4 = TFLongformerSelfAttention._mask_invalid_locations(hidden_states[:, :, 2:, :], 2)

        self.assertTrue(tf.math.reduce_sum(tf.cast(tf.math.is_inf(hid_states_1), tf.int64)) == 8)
        self.assertTrue(tf.math.reduce_sum(tf.cast(tf.math.is_inf(hid_states_2), tf.int64)) == 24)
        self.assertTrue(tf.math.reduce_sum(tf.cast(tf.math.is_inf(hid_states_3), tf.int64)) == 24)
        self.assertTrue(tf.math.reduce_sum(tf.cast(tf.math.is_inf(hid_states_4), tf.int64)) == 12)
//  [python]
def get_reprojection_loss(reconstruction_loss,multiScale=False, logs=False, weights=None,reduced=True):
	"""
	Build a lmbda op to be used to compute a loss function using reprojection between left and right frame
	Args:
		reconstruction_loss: name of the loss function used to compare reprojected and real image
		multiScale: if True compute multiple loss, one for each scale at which disparities are predicted
		logs: if True enable tf summary
		weights: array of weights to be multiplied for the losses at different resolution
		reduced: if true return the sum of the loss across the different scales, false to return an array with the different losses
	"""
	if reconstruction_loss not in ALL_LOSSES.keys():
		print('Unrecognized loss function, pick one among: {}'.format(ALL_LOSSES.keys()))
		raise Exception('Unknown loss function selected')
	base_loss_function = ALL_LOSSES[reconstruction_loss]
	if weights is None:
		weights = [1]*10
	def compute_loss(disparities,inputs):
		left = inputs['left']
		right = inputs['right']
		#normalize image to be between 0 and 1 
		left = tf.cast(left,dtype=tf.float32)/256.0
		right = tf.cast(right,dtype=tf.float32)/256.0
		accumulator=[]
		if multiScale:
			disp_to_test=len(disparities)
		else:
			disp_to_test=1
		for i in range(disp_to_test):
			#rescale prediction to full resolution
			current_disp = disparities[-(i+1)]
			disparity_scale_factor = tf.cast(tf.shape(left)[2],tf.float32)/tf.cast(tf.shape(current_disp)[2],tf.float32)
			resized_disp = preprocessing.resize_to_prediction(current_disp, left) * disparity_scale_factor

			reprojected_left = preprocessing.warp_image(right, resized_disp)
			partial_loss = base_loss_function(reprojected_left,left)
			if logs:
				tf.summary.scalar('Loss_resolution_{}'.format(i),partial_loss)
			accumulator.append(weights[i]*partial_loss)
		if reduced:
			return tf.reduce_sum(accumulator)
		else:
			return accumulator
	return compute_loss
//  [python]
def get_reduced_length(
    length,
    reduction_factor,
):
    return t.cast(
        t.ceil(t.div(length, t.cast(reduction_factor, dtype=length.dtype))),
        dtype=t.int32,
    )
//  [python]
def get_reduced_length(
    length,
    reduction_factor,
):
    return tf.cast(
        tf.math.ceil(tf.divide(length, tf.cast(reduction_factor, dtype=length.dtype))),
        dtype=tf.int32,
    )
//  [python]
def flow_to_warp(flow):
  """Compute the warp from the flow field.

  Args:
    flow: tf.tensor representing optical flow.

  Returns:
    The warp, i.e. the endpoints of the estimated flow.
  """

  # Construct a grid of the image coordinates.
  height, width = flow.shape.as_list()[-3:-1]
  i_grid, j_grid = tf.meshgrid(
      tf.linspace(0.0, height - 1.0, int(height)),
      tf.linspace(0.0, width - 1.0, int(width)),
      indexing='ij')
  grid = tf.stack([i_grid, j_grid], axis=2)

  # Potentially add batch dimension to match the shape of flow.
  if len(flow.shape) == 4:
    grid = grid[None]

  # Add the flow field to the image grid.
  if flow.dtype != grid.dtype:
    grid = t.cast(grid, flow.dtype)
  warp = grid + flow
  return warp
//  [python]
def flow_to_warp(flow):
  """Compute the warp from the flow field.

  Args:
    flow: tf.tensor representing optical flow.

  Returns:
    The warp, i.e. the endpoints of the estimated flow.
  """

  # Construct a grid of the image coordinates.
  height, width = flow.shape.as_list()[-3:-1]
  i_grid, j_grid = tf.meshgrid(
      tf.linspace(0.0, height - 1.0, int(height)),
      tf.linspace(0.0, width - 1.0, int(width)),
      indexing='ij')
  grid = tf.stack([i_grid, j_grid], axis=2)

  # Potentially add batch dimension to match the shape of flow.
  if len(flow.shape) == 4:
    grid = grid[None]

  # Add the flow field to the image grid.
  if flow.dtype != grid.dtype:
    grid = tf.cast(grid, flow.dtype)
  warp = grid + flow
  return warp
//  [python]
def mask_invalid(coords):
  """Mask coordinates outside of the image.

  Valid = 1, invalid = 0.

  Args:
    coords: a 4D float tensor of image coordinates.

  Returns:
    The mask showing which coordinates are valid.
  """
  coords_rank = len(coords.shape)
  if coords_rank != 4:
    raise NotImplementedError()
  max_height = float(coords.shape[-3] - 1)
  max_width = float(coords.shape[-2] - 1)
  mask = tf.logical_and(
      tf.logical_and(coords[:, :, :, 0] >= 0.0,
                     coords[:, :, :, 0] <= max_height),
      tf.logical_and(coords[:, :, :, 1] >= 0.0,
                     coords[:, :, :, 1] <= max_width))
  mask = tf.to_float(mask)[:, :, :, None]
  return mask
//  [python]
def mask_invalid(coords):
  """Mask coordinates outside of the image.

  Valid = 1, invalid = 0.

  Args:
    coords: a 4D float tensor of image coordinates.

  Returns:
    The mask showing which coordinates are valid.
  """
  coords_rank = len(coords.shape)
  if coords_rank != 4:
    raise NotImplementedError()
  max_height = float(coords.shape[-3] - 1)
  max_width = float(coords.shape[-2] - 1)
  mask = tf.logical_and(
      tf.logical_and(coords[:, :, :, 0] >= 0.0,
                     coords[:, :, :, 0] <= max_height),
      tf.logical_and(coords[:, :, :, 1] >= 0.0,
                     coords[:, :, :, 1] <= max_width))
  mask = tf.cast(mask, dtype=tf.float32)[:, :, :, None]
  return mask
//  [python]
def upsample(img, is_flow):
  """Double resolution of an image or flow field.

  Args:
    img: tf.tensor, image or flow field to be resized
    is_flow: bool, flag for scaling flow accordingly

  Returns:
    Resized and potentially scaled image or flow field.
  """
  _, height, width, _ = img.shape.as_list()
  orig_dtype = img.dtype
  if orig_dtype != tf.float32:
    img = tf.to_float(img)
  img_resized = tf.compat.v2.image.resize(img,
                                          (int(height * 2), int(width * 2)))
  if is_flow:
    # Scale flow values to be consistent with the new image size.
    img_resized *= 2
  if img_resized.dtype != orig_dtype:
    return tf.to_int32(img_resized)
  return img_resized
//  [python]
def upsample(img, is_flow):
  """Double resolution of an image or flow field.

  Args:
    img: tf.tensor, image or flow field to be resized
    is_flow: bool, flag for scaling flow accordingly

  Returns:
    Resized and potentially scaled image or flow field.
  """
  _, height, width, _ = img.shape.as_list()
  orig_dtype = img.dtype
  if orig_dtype != tf.float32:
    img = tf.cast(img, tf.float32)
  img_resized = tf.compat.v2.image.resize(img,
                                          (int(height * 2), int(width * 2)))
  if is_flow:
    # Scale flow values to be consistent with the new image size.
    img_resized *= 2
  if img_resized.dtype != orig_dtype:
    return tf.cast(img_resized, orig_dtype)
  return img_resized
//  [python]
def resample(source, coords):
  """Resample the source image at the passed coordinates.

  Args:
    source: tf.tensor, batch of images to be resampled.
    coords: tf.tensor, batch of coordinates in the image.

  Returns:
    The resampled image.

  Coordinates should be between 0 and size-1. Coordinates outside of this range
  are handled by interpolating with a background image filled with zeros in the
  same way that SAME size convolution works.
  """

  # Wrap this function because it uses a different order of height/width dims.
  orig_source_dtype = source.dtype
  if source.dtype != tf.float32:
    source = t.cast(source, tf.float32)
  if coords.dtype != tf.float32:
    coords = t.cast(coords, tf.float32)
  coords_rank = len(coords.shape)
  if coords_rank == 4:
    output = resampler(source, coords[:, :, :, ::-1])
    if orig_source_dtype != source.dtype:
      return t.cast(output, orig_source_dtype)
    return output
  else:
    raise NotImplementedError()
//  [python]
def resample(source, coords):
  """Resample the source image at the passed coordinates.

  Args:
    source: tf.tensor, batch of images to be resampled.
    coords: tf.tensor, batch of coordinates in the image.

  Returns:
    The resampled image.

  Coordinates should be between 0 and size-1. Coordinates outside of this range
  are handled by interpolating with a background image filled with zeros in the
  same way that SAME size convolution works.
  """

  # Wrap this function because it uses a different order of height/width dims.
  orig_source_dtype = source.dtype
  if source.dtype != tf.float32:
    source = tf.cast(source, tf.float32)
  if coords.dtype != tf.float32:
    coords = tf.cast(coords, tf.float32)
  coords_rank = len(coords.shape)
  if coords_rank == 4:
    output = resampler(source, coords[:, :, :, ::-1])
    if orig_source_dtype != source.dtype:
      return tf.cast(output, orig_source_dtype)
    return output
  else:
    raise NotImplementedError()
//  [python]
def masked_fill(
    tensor,
    mask,
    value=0,
):
    shape = shape_util.shape_list(tensor)
    mask = tf.broadcast_to(mask, shape)
    values = tf.fill(shape, value)
    return tf.where(mask, tensor, values)
//  [python]
def masked_fill(
    tensor,
    mask,
    value=0,
):
    shape = shape_util.shape_list(tensor)
    mask = tf.broadcast_to(mask, shape)
    values = tf.cast(tf.fill(shape, value), tensor.dtype)
    return tf.where(mask, tensor, values)
//  [python]
def bfloat16_to_float16(
    tensor,
):
    return tf.cond(tf.equal(tensor.dtype, tf.bfloat16), lambda: tf.cast(tensor, tf.float16), lambda: tensor)
//  [python]
def compute_time_length(
    tensor: tf.Tensor,
    dtype=tf.int32,
):
    with tf.name_scope("compute_time_length"):
        batch_size, time_length, *_ = shape_util.shape_list(tensor)
        return tf.cast(tf.repeat(time_length, batch_size, axis=0), dtype=dtype)
//  [python]
def apply_mask(
    outputs,
    mask=None,
    multiply=False,
):
    if mask is not None:
        expanded_mask = mask
        for _ in range(len(outputs.shape) - len(mask.shape)):  # expand last axis of mask so that it's dim equals output's dim
            expanded_mask = tf.expand_dims(expanded_mask, -1)
        if multiply:
            outputs = outputs * tf.cast(expanded_mask, dtype=outputs.dtype)
        outputs._keras_mask = mask  # pylint: disable=protected-access
    return outputs
//  [python]
def accuracy_ratio(logits_or_probs,labels):
    with tf.name_scope("accuracy_ratio"):
        _,actually_get = tf.nn.top_k(logits_or_probs,1)
        dim = len(actually_get.get_shape().as_list())-1
        actually_get = tf.squeeze(actually_get,axis=[dim])
        if labels.dtype != actually_get.dtype:
            actually_get = tf.to_int32(actually_get)
        correct_num = tf.reduce_sum(tf.to_float(tf.equal(actually_get,labels)))
        all_num = tf.to_float(wmlt.num_elements(labels))
        return correct_num*100./all_num
//  [python]
def accuracy_ratio(logits_or_probs,labels):
    with tf.name_scope("accuracy_ratio"):
        _,actually_get = tf.nn.top_k(logits_or_probs,1)
        dim = len(actually_get.get_shape().as_list())-1
        actually_get = tf.squeeze(actually_get,axis=[dim])
        if labels.dtype != actually_get.dtype:
            actually_get = tf.cast(actually_get,labels.dtype)
        correct_num = tf.reduce_sum(tf.cast(tf.equal(actually_get,labels),tf.float32))
        all_num = tf.cast(wmlt.num_elements(labels),tf.float32)
        return correct_num*100./all_num
//  [python]
def accuracy_num(logits,labels):
    with tf.name_scope("accuracy_num"):
        _,actually_get = tf.nn.top_k(logits,1)
        dim = len(actually_get.get_shape().as_list())-1
        actually_get = tf.squeeze(actually_get,axis=[dim])
        if labels.dtype != actually_get.dtype:
            actually_get = tf.to_int32(actually_get)
        correct_num = tf.reduce_sum(tf.to_float(tf.equal(actually_get,labels)))

    return correct_num,tf.reduce_prod(tf.shape(labels))
//  [python]
def accuracy_num(logits,labels):
    with tf.name_scope("accuracy_num"):
        _,actually_get = tf.nn.top_k(logits,1)
        dim = len(actually_get.get_shape().as_list())-1
        actually_get = tf.squeeze(actually_get,axis=[dim])
        if labels.dtype != actually_get.dtype:
            actually_get = tf.cast(actually_get,labels.dtype)
        correct_num = tf.reduce_sum(tf.cast(tf.equal(actually_get,labels),tf.float32))

    return correct_num,tf.reduce_prod(tf.shape(labels))
//  [python]
def sigmoid_accuracy_ratio(logits,labels,threshold=0.5):
    with tf.name_scope("accuracy_ratio"):
        assert len(logits.get_shape())==len(labels.get_shape()),"Error shape"
        prob = tf.sigmoid(logits)
        actually_get = prob>threshold
        actually_get = tf.to_float(actually_get)
        correct_num = tf.reduce_sum(tf.to_float(tf.equal(actually_get,labels)))
        all_num = tf.to_float(wmlt.num_elements(labels))
        return correct_num*100./all_num
//  [python]
def sigmoid_accuracy_ratio(logits,labels,threshold=0.5):
    with tf.name_scope("accuracy_ratio"):
        assert len(logits.get_shape())==len(labels.get_shape()),"Error shape"
        prob = tf.sigmoid(logits)
        actually_get = prob>threshold
        actually_get = tf.cast(actually_get,labels.dtype)
        correct_num = tf.reduce_sum(tf.cast(tf.equal(actually_get,labels),tf.float32))
        all_num = tf.cast(wmlt.num_elements(labels),tf.float32)
        return correct_num*100./all_num
//  [python]
def iou_of_multi_label(logits,labels,threshold=0.5):
    expected = tf.greater(labels,threshold)
    predicted = tf.greater(tf.nn.sigmoid(logits),threshold)
    union = tf.logical_or(expected,predicted)
    union = tf.reduce_sum(tf.to_float(union),axis=1)
    intersection = tf.logical_and(expected,predicted)
    intersection = tf.reduce_sum(tf.to_float(intersection),axis=1)
    def fn(u,i):
        return tf.cond(tf.greater(u,0.5),lambda:i/u,lambda:1.0)
    eval_results = tf.map_fn(lambda x:fn(x[0],x[1]),(union,intersection),dtype=tf.float32,back_prop=False)
    return eval_results
//  [python]
def iou_of_multi_label(logits,labels,threshold=0.5):
    expected = tf.greater(labels,threshold)
    predicted = tf.greater(tf.nn.sigmoid(logits),threshold)
    union = tf.logical_or(expected,predicted)
    union = tf.reduce_sum(tf.cast(union,tf.float32),axis=1)
    intersection = tf.logical_and(expected,predicted)
    intersection = tf.reduce_sum(tf.cast(intersection,tf.float32),axis=1)
    def fn(u,i):
        return tf.cond(tf.greater(u,0.5),lambda:i/u,lambda:1.0)
    eval_results = tf.map_fn(lambda x:fn(x[0],x[1]),(union,intersection),dtype=tf.float32,back_prop=False)
    return eval_results
//  [python]
def sigmoid_accuracy_ratiov2(logits,labels,threshold=0.5):
    with tf.name_scope("accuracy_ratio"):
        assert len(logits.get_shape())==len(labels.get_shape()),"Error shape"
        prob = tf.sigmoid(logits)
        if prob.get_shape().ndims>1:
            classes_nr = prob.get_shape().as_list()[1]
        if labels.dtype==tf.float32:
            labels = labels>threshold
            labels = tf.to_int32(labels)
        batch_size = tf.to_float(tf.shape(prob)[0])
        actually_get = prob>threshold
        actually_get = tf.to_int32(actually_get)
        if prob.get_shape().ndims>1:
            correct_num = tf.reduce_sum(tf.to_float(tf.equal(actually_get,labels)),axis=1)
            correct_num = tf.reduce_sum(tf.to_float(tf.equal(correct_num,classes_nr)))
        else:
            correct_num = tf.reduce_sum(tf.to_float(tf.equal(actually_get,labels)))

        return correct_num*100./batch_size
//  [python]
def sigmoid_accuracy_ratiov2(logits,labels,threshold=0.5):
    with tf.name_scope("accuracy_ratio"):
        assert len(logits.get_shape())==len(labels.get_shape()),"Error shape"
        prob = tf.sigmoid(logits)
        if prob.get_shape().ndims>1:
            classes_nr = prob.get_shape().as_list()[1]
        if labels.dtype==tf.float32:
            labels = labels>threshold
            labels = tf.cast(labels,tf.int32)
        batch_size = tf.cast(tf.shape(prob)[0],tf.float32)
        actually_get = prob>threshold
        actually_get = tf.cast(actually_get,labels.dtype)
        if prob.get_shape().ndims>1:
            correct_num = tf.reduce_sum(tf.cast(tf.equal(actually_get,labels),tf.float32),axis=1)
            correct_num = tf.reduce_sum(tf.cast(tf.equal(correct_num,classes_nr),tf.float32))
        else:
            correct_num = tf.reduce_sum(tf.cast(tf.equal(actually_get,labels),tf.float32))

        return correct_num*100./batch_size
//  [python]
def random_apply(func, p, x):
  """Randomly apply function func to x with probability p."""
  return tf.cond(
      tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
              tf.cast(p, tf.float32)),
      lambda: func(x),
      lambda: x)
//  [python]
def random_solarization(image, p=0.2):
  """Random solarize image."""
  def _transform(image):
    image = image * t.cast(t.less(image, 0.5), dtype=image.dtype) + (
        1.0 - image) * t.cast(t.greater_equal(image, 0.5), dtype=image.dtype)
    return image
  return random_apply(_transform, p=p, x=image)
//  [python]
def random_solarization(image, p=0.2):
  """Random solarize image."""
  def _transform(image):
    image = image * tf.cast(tf.less(image, 0.5), dtype=image.dtype) + (
        1.0 - image) * tf.cast(tf.greater_equal(image, 0.5), dtype=image.dtype)
    return image
  return random_apply(_transform, p=p, x=image)
//  [python]
def _train_step(self, data):
        images, targets = data

        with tf.GradientTape() as tape:
            predictions = self._model(images, training=True)
            loss = self._model.loss(targets, predictions)
            loss['total-loss'] = loss['weighted-loss']

            if self.params.training.use_weight_decay:
                loss['l2-regularization'] = self.weight_decay()
                loss['total-loss'] += loss['l2-regularization']

            per_replica_loss = loss['total-loss'] / self.num_replicas

            if self.use_float16:
                per_replica_loss = self.optimizer.get_scaled_loss(
                    per_replica_loss)

        gradients = tape.gradient(per_replica_loss,
                                  self._model.trainable_variables)
        if self.use_float16:
            gradients = self.optimizer.get_unscaled_gradients(gradients)

        if self._clip_gradients:
            threshold = self.params.training.optimizer.clipnorm
            gradients = Executor._clip_gradients_by_norm(gradients, threshold)

        self.optimizer.apply_gradients(
            zip(gradients, self._model.trainable_variables))

        loss['num-anchors-matched'] /= tf.to_float(tf.shape(images)[0])
        loss['gradient-norm'] = tf.linalg.global_norm(gradients) * self.num_replicas
        return loss
//  [python]
def _train_step(self, data):
        images, targets = data

        with tf.GradientTape() as tape:
            predictions = self._model(images, training=True)
            loss = self._model.loss(targets, predictions)
            loss['total-loss'] = loss['weighted-loss']

            if self.params.training.use_weight_decay:
                loss['l2-regularization'] = self.weight_decay()
                loss['total-loss'] += loss['l2-regularization']

            per_replica_loss = loss['total-loss'] / self.num_replicas

            if self.use_float16:
                per_replica_loss = self.optimizer.get_scaled_loss(
                    per_replica_loss)

        gradients = tape.gradient(per_replica_loss,
                                  self._model.trainable_variables)
        if self.use_float16:
            gradients = self.optimizer.get_unscaled_gradients(gradients)

        if self._clip_gradients:
            threshold = self.params.training.optimizer.clipnorm
            gradients = Executor._clip_gradients_by_norm(gradients, threshold)

        self.optimizer.apply_gradients(
            zip(gradients, self._model.trainable_variables))

        loss['num-anchors-matched'] /= tf.cast(tf.shape(images)[0], dtype=tf.float32)
        loss['gradient-norm'] = tf.linalg.global_norm(gradients) * self.num_replicas
        return loss
//  [python]
def sparse_softmax_cross_entropy_with_logits_alpha_balanced(
        _sentinel=None,  # pylint: disable=invalid-name
        labels=None,
        logits=None,
        alpha="auto",
        max_alpha_scale=10.0,
        name=None):
    if alpha is None:
        return tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels,logits=logits,name=name)
    elif alpha != "auto":
        raise ValueError(f"alphpa shold be None or auto, not {alpha}.")
    with tf.variable_scope(name,default_name="sparse_softmax_cross_entropy_with_logits_alpha_balanced"):
        probability = tf.nn.softmax(logits)
        labels = tf.expand_dims(labels,axis=-1)
        r_probability = wmlt.batch_gather(probability,labels)
        r_probability = tf.squeeze(r_probability,axis=-1)
        r_probability = tf.maximum(1e-10*(1+r_probability),r_probability)
        num_classes = logits.get_shape().as_list()[-1]
        labels = tf.squeeze(labels,axis=-1)
        one_hot_labels = tf.one_hot(indices=labels,depth=num_classes)
        count = tf.reduce_sum(one_hot_labels,axis=list(range(one_hot_labels.get_shape().ndims-1)),keepdims=True)
        count = tf.maximum(count,1)
        ratio = tf.reduce_prod(tf.shape(labels))/(count*float(num_classes))
        alpha_t = tf.minimum(max_alpha_scale,1.0/ratio)
        alpha_t = alpha_t*tf.ones_like(one_hot_labels,dtype=tf.float32)
        labels = tf.expand_dims(labels,axis=-1)
        alpha_t = wmlt.batch_gather(params=alpha_t,indices=labels)
        alpha_t = tf.squeeze(alpha_t,axis=-1)
        loss = -alpha_t*tf.math.log(r_probability)
        return loss