jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_new_alive_state(self, new_seq: tf.Tensor, new_log_probs: tf.Tensor,
                           new_finished_flags: tf.Tensor,
                           new_cache: Dict[str, tf.Tensor]) -> Dict[str, Any]:
    """Gather the sequences that are still alive.

    This function resets the sequences in the alive_state that are finished.

    Args:
      new_seq: New sequences generated by growing the current alive sequences
        int32 tensor with shape [batch_size, cur_index + 1]
      new_log_probs: Log probabilities of new sequences float32 tensor with
        shape [batch_size, 1]
      new_finished_flags: A boolean Tensor indicates which sequences are live
        inside the beam.
      new_cache: Dict of cached values for each sequence.

    Returns:
      Dictionary with alive keys.
    """
    new_seq = tf.multiply(
        new_seq, tf.to_int32(tf.logical_not(new_finished_flags)))
    return {
        decoding_module.StateKeys.ALIVE_SEQ: new_seq,
        decoding_module.StateKeys.ALIVE_LOG_PROBS: new_log_probs,
        decoding_module.StateKeys.ALIVE_CACHE: new_cache
    }
//  [python]
def _get_new_alive_state(self, new_seq: tf.Tensor, new_log_probs: tf.Tensor,
                           new_finished_flags: tf.Tensor,
                           new_cache: Dict[str, tf.Tensor]) -> Dict[str, Any]:
    """Gather the sequences that are still alive.

    This function resets the sequences in the alive_state that are finished.

    Args:
      new_seq: New sequences generated by growing the current alive sequences
        int32 tensor with shape [batch_size, cur_index + 1]
      new_log_probs: Log probabilities of new sequences float32 tensor with
        shape [batch_size, 1]
      new_finished_flags: A boolean Tensor indicates which sequences are live
        inside the beam.
      new_cache: Dict of cached values for each sequence.

    Returns:
      Dictionary with alive keys.
    """
    new_seq = tf.multiply(
        new_seq, tf.cast(tf.logical_not(new_finished_flags), new_seq.dtype))
    return {
        decoding_module.StateKeys.ALIVE_SEQ: new_seq,
        decoding_module.StateKeys.ALIVE_LOG_PROBS: new_log_probs,
        decoding_module.StateKeys.ALIVE_CACHE: new_cache
    }
//  [python]
def _evaluate(self, states, actions, tasks):
    # X velocity is the first component of the velocity vector.
    vel = states[:, self._vel_index]
    target_vel = tf.squeeze(tasks, axis=1)
    if self._margin is None:
      vel_rew = -1.0 * tf.abs(target_vel - vel)
    else:
      vel_rew = tf.cast(tf.abs(target_vel - vel) < self._margin, tf.float32)
    ctrl_cost = 0.05 * tf.reduce_sum(actions**2, axis=1)
    rew = vel_rew - ctrl_cost
    dones = tf.cast(states[:, -1], tf.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def quadruped_reward(states,
                     actions,
                     goals,
                     sparse=False,
                     goal_distance=None,
                     goal_radius=1.0):
  """Reward function for the quadruped.

  Args:
    states: the current state.
    actions: the current action.
    goals: the goal state.
    sparse: whether the reward function is sparse or dense.
    goal_distance: *initial* distance to the goal.
    goal_radius: *desired* distance to the goal.
  Returns:
    reward: the reward.
  """
  assert goal_distance is not None
  pos = states[:, :2]
  dist_to_goal = tf.linalg.norm(pos - goals, axis=1)
  ctrl_cost = 0.1 * tf.reduce_sum(actions**2, axis=1)
  stay_alive_bonus = 1.0
  if sparse:
    rew = t.cast(dist_to_goal < goal_radius, t.float32) - 1.0
    dones = t.cast(tf.zeros_like(rew), t.bool)
  else:
    rew = (goal_distance - dist_to_goal) - ctrl_cost + stay_alive_bonus
    dones = t.cast(states[:, -1], t.bool)
  assert rew.shape == dones.shape
  return rew, dones
//  [python]
def quadruped_reward(states,
                     actions,
                     goals,
                     sparse=False,
                     goal_distance=None,
                     goal_radius=1.0):
  """Reward function for the quadruped.

  Args:
    states: the current state.
    actions: the current action.
    goals: the goal state.
    sparse: whether the reward function is sparse or dense.
    goal_distance: *initial* distance to the goal.
    goal_radius: *desired* distance to the goal.
  Returns:
    reward: the reward.
  """
  assert goal_distance is not None
  pos = states[:, :2]
  dist_to_goal = tf.linalg.norm(pos - goals, axis=1)
  ctrl_cost = 0.1 * tf.reduce_sum(actions**2, axis=1)
  stay_alive_bonus = 1.0
  if sparse:
    rew = tf.cast(dist_to_goal < goal_radius, tf.float32) - 1.0
    dones = tf.cast(tf.zeros_like(rew), tf.bool)
  else:
    rew = (goal_distance - dist_to_goal) - ctrl_cost + stay_alive_bonus
    dones = tf.cast(states[:, -1], tf.bool)
  assert rew.shape == dones.shape
  return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    height = states[:, self._height_index]
    x_vel = states[:, self._vel_index + 1]
    left_foot_index = 2 * 3 + self._orientation_index
    right_foot_index = 2 * 6 + self._orientation_index
    xz_left_foot = states[:, left_foot_index:left_foot_index + 2]
    xz_right_foot = states[:, right_foot_index:right_foot_index + 2]
    ctrl = tf.reduce_sum(actions**2, axis=1)
    height_offset = 0.5
    foot_offset = np.array([0.0, -0.7])[None]
    features = tf.concat(
        [
            height[:, None] - height_offset,
            self._vel_coef * x_vel[:, None],
            xz_left_foot - foot_offset,
            xz_right_foot - foot_offset,
            ctrl[:, None],
        ],
        axis=1,
    )
    assert features.shape == tasks.shape
    rew = tf.reduce_sum(features * tasks, axis=1) + self._stay_alive_bonus
    if self._terminate_on_fall:  # DM control uses 1.2 as the stand height
      dones = height < self._fall_height
    else:
      dones = t.cast(states[:, -1], t.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    height = states[:, self._height_index]
    x_vel = states[:, self._vel_index + 1]
    left_foot_index = 2 * 3 + self._orientation_index
    right_foot_index = 2 * 6 + self._orientation_index
    xz_left_foot = states[:, left_foot_index:left_foot_index + 2]
    xz_right_foot = states[:, right_foot_index:right_foot_index + 2]
    ctrl = tf.reduce_sum(actions**2, axis=1)
    height_offset = 0.5
    foot_offset = np.array([0.0, -0.7])[None]
    features = tf.concat(
        [
            height[:, None] - height_offset,
            self._vel_coef * x_vel[:, None],
            xz_left_foot - foot_offset,
            xz_right_foot - foot_offset,
            ctrl[:, None],
        ],
        axis=1,
    )
    assert features.shape == tasks.shape
    rew = tf.reduce_sum(features * tasks, axis=1) + self._stay_alive_bonus
    if self._terminate_on_fall:  # DM control uses 1.2 as the stand height
      dones = height < self._fall_height
    else:
      dones = tf.cast(states[:, -1], tf.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _get_new_finished_state(self, state: Dict[str, Any], new_seq: tf.Tensor,
                              new_log_probs: tf.Tensor,
                              new_finished_flags: tf.Tensor,
                              batch_size: int) -> Dict[str, tf.Tensor]:
    """Combine new and old finished sequences.

    Args:
      state: A dictionary with the current loop state.
      new_seq: New sequences generated by growing the current alive sequences
        int32 tensor [batch, curr_index + 1] or [batch, max_decode_length + 1].
      new_log_probs: Log probabilities of new sequences float32 tensor with
        shape [batch, 1].
      new_finished_flags: A boolean Tensor indicates which sequences are live.
      batch_size: The given batch size.

    Returns:
      Dictionary with finished keys from StateKeys.
    """
    i = state[decoding_module.StateKeys.CUR_INDEX]
    finished_seq = state[decoding_module.StateKeys.FINISHED_SEQ]
    finished_scores = state[decoding_module.StateKeys.FINISHED_SCORES]
    finished_flags = state[decoding_module.StateKeys.FINISHED_FLAGS]

    if not self.padded_decode:
      finished_seq = tf.concat(
          [finished_seq, tf.zeros([batch_size, 1], tf.int32)], axis=-1)
    new_scores = new_log_probs
    if self.length_normalization_fn is not None:
      length_norm = self.length_normalization_fn(i + 1, self.dtype)
      new_scores = new_log_probs / length_norm
    new_seq = tf.multiply(
        new_seq, tf.to_float(tf.logical_not(finished_flags)))
    new_scores = tf.multiply(
        new_scores, tf.to_float(tf.logical_not(finished_flags)))

    finished_seq += tf.multiply(new_seq,
                                tf.to_float(new_finished_flags))
    finished_scores += tf.multiply(
        new_scores, tf.to_float(new_finished_flags))
    new_finished_flags = tf.logical_or(new_finished_flags, finished_flags)
    return {
        decoding_module.StateKeys.FINISHED_SEQ: finished_seq,
        decoding_module.StateKeys.FINISHED_SCORES: finished_scores,
        decoding_module.StateKeys.FINISHED_FLAGS: new_finished_flags
    }
//  [python]
def _get_new_finished_state(self, state: Dict[str, Any], new_seq: tf.Tensor,
                              new_log_probs: tf.Tensor,
                              new_finished_flags: tf.Tensor,
                              batch_size: int) -> Dict[str, tf.Tensor]:
    """Combine new and old finished sequences.

    Args:
      state: A dictionary with the current loop state.
      new_seq: New sequences generated by growing the current alive sequences
        int32 tensor [batch, curr_index + 1] or [batch, max_decode_length + 1].
      new_log_probs: Log probabilities of new sequences float32 tensor with
        shape [batch, 1].
      new_finished_flags: A boolean Tensor indicates which sequences are live.
      batch_size: The given batch size.

    Returns:
      Dictionary with finished keys from StateKeys.
    """
    i = state[decoding_module.StateKeys.CUR_INDEX]
    finished_seq = state[decoding_module.StateKeys.FINISHED_SEQ]
    finished_scores = state[decoding_module.StateKeys.FINISHED_SCORES]
    finished_flags = state[decoding_module.StateKeys.FINISHED_FLAGS]

    if not self.padded_decode:
      finished_seq = tf.concat(
          [finished_seq, tf.zeros([batch_size, 1], tf.int32)], axis=-1)
    new_scores = new_log_probs
    if self.length_normalization_fn is not None:
      length_norm = self.length_normalization_fn(i + 1, self.dtype)
      new_scores = new_log_probs / length_norm
    new_seq = tf.multiply(
        new_seq, tf.cast(tf.logical_not(finished_flags), new_seq.dtype))
    new_scores = tf.multiply(
        new_scores, tf.cast(tf.logical_not(finished_flags), new_scores.dtype))

    finished_seq += tf.multiply(new_seq,
                                tf.cast(new_finished_flags, new_seq.dtype))
    finished_scores += tf.multiply(
        new_scores, tf.cast(new_finished_flags, new_scores.dtype))
    new_finished_flags = tf.logical_or(new_finished_flags, finished_flags)
    return {
        decoding_module.StateKeys.FINISHED_SEQ: finished_seq,
        decoding_module.StateKeys.FINISHED_SCORES: finished_scores,
        decoding_module.StateKeys.FINISHED_FLAGS: new_finished_flags
    }
//  [python]
def generate_pair(self, key):
    """Samples sequence pair from (simplified) pair HMM given seed."""
    states = self.sample_states(key)
    aln_len = tf.shape(states)[0]
    tokens = self._sampler.sample([aln_len])

    em_x = t.cast(
        tf.logical_or(states == self.MATCH, states == self.GAP_IN_Y), t.int32)
    em_y = t.cast(
        tf.logical_or(states == self.MATCH, states == self.GAP_IN_X), t.int32)
    seq_x = tf.gather(tokens, tf.squeeze(tf.where(em_x), axis=1))
    seq_y = tf.gather(tokens, tf.squeeze(tf.where(em_y), axis=1))
    key_x, key_y = split_key(key)
    seq_x, start_x = self.extend_tails(key_x, seq_x)
    seq_y, start_y = self.extend_tails(key_y, seq_y)

    pos_x = start_x + tf.cumsum(em_x)
    pos_y = start_y + tf.cumsum(em_y)
    enc_trans = tf.gather(self._trans_encoder,
                          self._hash_fn(states[:-1], states[1:]))
    enc_trans = tf.concat([self._init_trans, enc_trans], 0)
    alignment_targets = tf.stack([pos_x, pos_y, enc_trans])

    sequence = tf.stack([seq_x, seq_y])
    aln_to_pad = tf.maximum(0, 2 * self._max_len - aln_len)
    alignment_targets = tf.pad(alignment_targets, [[0, 0], [0, aln_to_pad]])

    fam_key = tf.random.stateless_uniform(
        [], key, maxval=tf.int32.max, dtype=tf.int32)
    fam_key = tf.repeat(fam_key, 2)
    sequence.set_shape([2, self._max_len])
    alignment_targets.set_shape([3, 2 * self._max_len])
    fam_key.set_shape([2])
    pid1_key, pid3_key = split_key(key)
    pid1 = tf.random.stateless_uniform([], pid1_key)
    pid3 = tf.random.stateless_uniform([], pid3_key)

    return {
        'sequence': sequence,
        'alignment/targets': alignment_targets,
        'fam_key': fam_key,
        'alignment/pid1': pid1,
        'alignment/pid3': pid3,
    }
//  [python]
def generate_pair(self, key):
    """Samples sequence pair from (simplified) pair HMM given seed."""
    states = self.sample_states(key)
    aln_len = tf.shape(states)[0]
    tokens = self._sampler.sample([aln_len])

    em_x = tf.cast(
        tf.logical_or(states == self.MATCH, states == self.GAP_IN_Y), tf.int32)
    em_y = tf.cast(
        tf.logical_or(states == self.MATCH, states == self.GAP_IN_X), tf.int32)
    seq_x = tf.gather(tokens, tf.squeeze(tf.where(em_x), axis=1))
    seq_y = tf.gather(tokens, tf.squeeze(tf.where(em_y), axis=1))
    key_x, key_y = split_key(key)
    seq_x, start_x = self.extend_tails(key_x, seq_x)
    seq_y, start_y = self.extend_tails(key_y, seq_y)

    pos_x = start_x + tf.cumsum(em_x)
    pos_y = start_y + tf.cumsum(em_y)
    enc_trans = tf.gather(self._trans_encoder,
                          self._hash_fn(states[:-1], states[1:]))
    enc_trans = tf.concat([self._init_trans, enc_trans], 0)
    alignment_targets = tf.stack([pos_x, pos_y, enc_trans])

    sequence = tf.stack([seq_x, seq_y])
    aln_to_pad = tf.maximum(0, 2 * self._max_len - aln_len)
    alignment_targets = tf.pad(alignment_targets, [[0, 0], [0, aln_to_pad]])

    fam_key = tf.random.stateless_uniform(
        [], key, maxval=tf.int32.max, dtype=tf.int32)
    fam_key = tf.repeat(fam_key, 2)
    sequence.set_shape([2, self._max_len])
    alignment_targets.set_shape([3, 2 * self._max_len])
    fam_key.set_shape([2])
    pid1_key, pid3_key = split_key(key)
    pid1 = tf.random.stateless_uniform([], pid1_key)
    pid3 = tf.random.stateless_uniform([], pid3_key)

    return {
        'sequence': sequence,
        'alignment/targets': alignment_targets,
        'fam_key': fam_key,
        'alignment/pid1': pid1,
        'alignment/pid3': pid3,
    }
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    tip_pos = self.state_to_task(states)
    target_tip_pos = tasks
    dist = tf.norm(tip_pos - target_tip_pos, axis=1)
    dones = dist < self._margin
    if self._use_neg_rew:
      rew = t.cast(dones, t.float32) - 1.0
    else:
      rew = t.cast(dones, t.float32)

    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    tip_pos = self.state_to_task(states)
    target_tip_pos = tasks
    dist = tf.norm(tip_pos - target_tip_pos, axis=1)
    dones = dist < self._margin
    if self._use_neg_rew:
      rew = tf.cast(dones, tf.float32) - 1.0
    else:
      rew = tf.cast(dones, tf.float32)

    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, self._pos_index:self._pos_index + 2]
    target_pos = tasks
    dist = tf.norm(pos - target_pos, axis=1)
    dones = dist < self._margin
    if self._use_neg_rew:
      rew = t.cast(dones, t.float32) - 1.0
    else:
      rew = t.cast(dones, t.float32)

    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, self._pos_index:self._pos_index + 2]
    target_pos = tasks
    dist = tf.norm(pos - target_pos, axis=1)
    dones = dist < self._margin
    if self._use_neg_rew:
      rew = tf.cast(dones, tf.float32) - 1.0
    else:
      rew = tf.cast(dones, tf.float32)

    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    height = states[:, self._head_height_index]
    extremities = states[:,
                         self._extremities_index:self._extremities_index + 12]
    extremities_dist = tf.norm(
        tf.reshape(extremities, (-1, 4, 3)),
        axis=-1)  # distance to the 4 extremities
    com_vel = states[:, self._com_vel_index:self._com_vel_index + 3]
    ctrl = tf.reduce_sum(actions**2, axis=1)
    height_offset = 1.0  # DM uses 1.4 as the standing threshold
    extremities_offset = 0.5
    com_vel_offset = 0.0
    features = tf.concat(
        [
            height[:, None] - height_offset,  # 1
            extremities_dist - extremities_offset,  # 4
            com_vel - com_vel_offset,  # 3
            ctrl[:, None],
        ],
        axis=1,
    )  # 1
    fall_penalty = -1.0 * self._fall_penalty * tf.to_float(height < 1.4)
    assert features.shape == tasks.shape
    rew = tf.reduce_sum(features * tasks, axis=1) + fall_penalty
    dones = tf.to_bool(states[:, -1])
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    height = states[:, self._head_height_index]
    extremities = states[:,
                         self._extremities_index:self._extremities_index + 12]
    extremities_dist = tf.norm(
        tf.reshape(extremities, (-1, 4, 3)),
        axis=-1)  # distance to the 4 extremities
    com_vel = states[:, self._com_vel_index:self._com_vel_index + 3]
    ctrl = tf.reduce_sum(actions**2, axis=1)
    height_offset = 1.0  # DM uses 1.4 as the standing threshold
    extremities_offset = 0.5
    com_vel_offset = 0.0
    features = tf.concat(
        [
            height[:, None] - height_offset,  # 1
            extremities_dist - extremities_offset,  # 4
            com_vel - com_vel_offset,  # 3
            ctrl[:, None],
        ],
        axis=1,
    )  # 1
    fall_penalty = -1.0 * self._fall_penalty * tf.cast(height < 1.4, tf.float32)
    assert features.shape == tasks.shape
    rew = tf.reduce_sum(features * tasks, axis=1) + fall_penalty
    dones = tf.cast(states[:, -1], tf.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, :2]
    target_pos = tasks
    dist = tf.norm(pos - target_pos, axis=1)
    target_size = 0.1
    rew = tf_tolerance(
        dist,
        bounds=(0, target_size),
        margin=5 * target_size,
        sigmoid="long_tail")
    dones = t.cast(states[:, -1], t.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, :2]
    target_pos = tasks
    dist = tf.norm(pos - target_pos, axis=1)
    target_size = 0.1
    rew = tf_tolerance(
        dist,
        bounds=(0, target_size),
        margin=5 * target_size,
        sigmoid="long_tail")
    dones = tf.cast(states[:, -1], tf.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def custom_scale_grad(x, num, clamp_value):
    def grad(dy):
        gs = 1. / tf.math.sqrt(tf.cast(num, tf.float32) * clamp_value + 0.0)
        return dy * gs, None, None
    
    return tf.identity(x), grad
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, self._pos_index:self._pos_index + 2]
    dist = tf.norm(
        pos[:, None] - self._goal_vec[None], axis=-1)  # B x num_tasks
    if self._linear:
      at_goals = t.cast(dist < self._margin, t.float32)
      rew = tf.reduce_sum(at_goals * tasks, axis=1)
      dones = t.cast(tf.zeros_like(rew), t.bool)
    else:
      dist = tf.reduce_sum(dist * tasks, axis=1)
      dones = dist < self._margin
      rew = t.cast(dones, t.float32) - 1.0
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, self._pos_index:self._pos_index + 2]
    dist = tf.norm(
        pos[:, None] - self._goal_vec[None], axis=-1)  # B x num_tasks
    if self._linear:
      at_goals = tf.cast(dist < self._margin, tf.float32)
      rew = tf.reduce_sum(at_goals * tasks, axis=1)
      dones = tf.cast(tf.zeros_like(rew), tf.bool)
    else:
      dist = tf.reduce_sum(dist * tasks, axis=1)
      dones = dist < self._margin
      rew = tf.cast(dones, tf.float32) - 1.0
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, :3]
    upright = states[:, self._upright_index]
    target_pos = tasks
    dist = tf.norm(pos - target_pos, axis=1)

    radii = 0.005 + 0.04
    is_upright = 0.5 * (upright + 0.5)
    in_target = tf_tolerance(
        dist, bounds=(0, radii), margin=2 * radii, sigmoid="gaussian")
    rew = (7 * in_target + is_upright) / 8
    dones = t.cast(states[:, -1], t.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def _evaluate(self, states, actions, tasks):
    del actions
    pos = states[:, :3]
    upright = states[:, self._upright_index]
    target_pos = tasks
    dist = tf.norm(pos - target_pos, axis=1)

    radii = 0.005 + 0.04
    is_upright = 0.5 * (upright + 0.5)
    in_target = tf_tolerance(
        dist, bounds=(0, radii), margin=2 * radii, sigmoid="gaussian")
    rew = (7 * in_target + is_upright) / 8
    dones = tf.cast(states[:, -1], tf.bool)
    assert rew.shape == dones.shape
    return rew, dones
//  [python]
def call(self, x):
        import tensorflow as tf
        return tf.one_hot(tf.to_int32(x), self.classes)
//  [python]
def call(self, x):
        import tensorflow as tf
        return tf.one_hot(tf.cast(x, tf.int32), self.classes)
//  [python]
def __call__(self, logits, samples, logit_length=None):
        labels = insert_eos_in_labels(samples["output"], self.eos, samples["output_length"])
        mask = tf.math.logical_not(tf.math.equal(labels, 0))
        labels = tf.one_hot(indices=labels, depth=self.num_classes)
        seq_len = tf.shape(labels)[1]
        logits = logits[:, :seq_len, :]
        loss = self.call(labels, logits)
        mask = tf.cast(mask, dtype=loss.dtype)
        loss *= mask
        if self.by_token:
            return tf.divide(tf.reduce_sum(loss), tf.reduce_sum(mask))
        if self.by_sequence:
            loss = tf.reduce_sum(loss, axis=-1)
        return tf.reduce_mean(loss)
//  [python]
def batch_gather():
    def mselect(args):
        import tensorflow as tf
        logits, seqi = args
        idx = tf.to_int32(seqi)
        logits = tf.batch_gather(logits, idx)
        return logits
    return mselect
//  [python]
def batch_gather():
    def mselect(args):
        import tensorflow as tf
        logits, seqi = args
        idx = tf.cast(seqi, tf.int32)
        logits = tf.batch_gather(logits, idx)
        return logits
    return mselect
//  [python]
def get_num_dims(x, /, *, as_array=False):
    return (
        t.cast(tf.shape(tf.shape(x))[0], tf.int64)
        if as_array
        else int(tf.shape(tf.shape(x)))
    )
//  [python]
def get_num_dims(x, /, *, as_array=False):
    return (
        tf.cast(tf.shape(tf.shape(x))[0], tf.int64)
        if as_array
        else int(tf.shape(tf.shape(x)))
    )
//  [python]
def SeqLikelihood(seq_model, history_length, name='', reverse=False, lf=5):
    from keras.layers import Lambda
    name = 'sequence_log_likelihood_' + name

    def bayes(seq):
        import tensorflow as tf
        seqi = seq
        if reverse:
            seqi = seq[:, -history_length - lf:]
        else:
            seqi = seq[:, : history_length + lf]

        logits = seq_model(seqi)
        # logits = tf.nn.softmax(logits)
        logits = tf.log(logits + 1e-18)
        idx = tf.cast(tf.expand_dims(seqi, axis=-1), tf.int32)
        if reverse:
            logits = tf.squeeze(tf.batch_gather(logits[:, 1:], idx[:, :-1]), axis=-1)
        else:
            logits = tf.squeeze(tf.batch_gather(logits[:, :-1], idx[:, 1:]), axis=-1)
        if reverse:
            logits = logits[:, : -history_length + 1]
        else:
            logits = logits[:, history_length - 1:]
        logits = tf.cumsum(logits, axis=-1, reverse=reverse)
        return logits
        """
        if reverse:
            logits = logits / tf.expand_dims(tf.range(10, 0, -1, dtype=tf.float32), axis=0)
        else:
            logits = logits / tf.expand_dims(tf.range(1, 6, dtype=tf.float32), axis=0)
        return idx  # logits
        """
    return Lambda(bayes, name=name, dtype='int32')
//  [python]
def generate_same_image_mask(num_pixels):
  """Generates a mask indicating if two pixels belong to the same image or not.

  Args:
    num_pixels: List with `num_image` elements representing the number of
      pixels in each image.

  Returns:
    A tensor of shape [1, num_total_pixels, num_total_pixels] indicating
      which pixel pairs belong to the same image.
  """
  image_ids = []
  num_total_pixels = 0
  for img_id, pixel_count in enumerate(num_pixels):
    image_ids += ([img_id] * pixel_count)
    num_total_pixels += pixel_count

  image_ids = tf.reshape(tf.constant(image_ids), shape=(num_total_pixels, 1))
  same_image_mask = tf.to_float(
      tf.equal(image_ids, tf.transpose(image_ids)))
  return tf.expand_dims(same_image_mask, axis=0)
//  [python]
def generate_same_image_mask(num_pixels):
  """Generates a mask indicating if two pixels belong to the same image or not.

  Args:
    num_pixels: List with `num_image` elements representing the number of
      pixels in each image.

  Returns:
    A tensor of shape [1, num_total_pixels, num_total_pixels] indicating
      which pixel pairs belong to the same image.
  """
  image_ids = []
  num_total_pixels = 0
  for img_id, pixel_count in enumerate(num_pixels):
    image_ids += ([img_id] * pixel_count)
    num_total_pixels += pixel_count

  image_ids = tf.reshape(tf.constant(image_ids), shape=(num_total_pixels, 1))
  same_image_mask = tf.cast(
      tf.equal(image_ids, tf.transpose(image_ids)), dtype=tf.float32)
  return tf.expand_dims(same_image_mask, axis=0)
//  [python]
main_graph = tf.Graph()
with main_graph.as_default():
  ids1 = tf.compat.v1.placeholder(dtype=tf.int32, name='ids1')
  ids2 = tf.compat.v1.placeholder(dtype=tf.int32, name='ids2')
  casted_ids1 = tf.to_float(ids1)
  casted_ids2 = tf.to_float(ids2)

  remote_a0 = remote_op_a(ids1)
  remote_b0 = remote_op_b(ids1, ids2)

  left_upper_concat = tf.concat([remote_a0, remote_b0], 0)
  left_upper_sum = tf.reduce_mean(left_upper_concat)

  right_upper_sum = tf.reduce_mean(remote_b0)
  right_upper_mul = tf.multiply(right_upper_sum, casted_ids2)
  right_upper_add = tf.add(right_upper_mul, left_upper_sum)
  right_upper_round = tf.math.round(right_upper_mul)
  right_upper_floormod = tf.math.floormod(right_upper_round, N)

  left_upper_add = tf.add_n([left_upper_sum, casted_ids1, right_upper_add])
  left_upper_round = tf.math.round(left_upper_add)
  left_upper_floormod = tf.math.floormod(left_upper_round, N)

  remote_a1 = remote_op_a(left_upper_floormod)
  remote_b1 = remote_op_b(left_upper_floormod, right_upper_floormod)

  left_lower_sum = tf.reduce_mean(remote_a1)

  right_lower_sum = tf.reduce_mean(remote_b1)
  right_lower_mul = tf.multiply(casted_ids2, right_lower_sum)
  right_lower_div = tf.divide(right_upper_add, right_lower_mul)

  main_result = tf.add_n([
      left_lower_sum, right_lower_div, right_lower_sum, right_upper_sum,
      tf.to_float(left_upper_floormod)
  ])
//  [python]
main_graph = tf.Graph()
with main_graph.as_default():
  ids1 = tf.compat.v1.placeholder(dtype=tf.int32, name='ids1')
  ids2 = tf.compat.v1.placeholder(dtype=tf.int32, name='ids2')
  casted_ids1 = tf.cast(ids1, tf.float32)
  casted_ids2 = tf.cast(ids2, tf.float32)

  remote_a0 = remote_op_a(ids1)
  remote_b0 = remote_op_b(ids1, ids2)

  left_upper_concat = tf.concat([remote_a0, remote_b0], 0)
  left_upper_sum = tf.reduce_mean(left_upper_concat)

  right_upper_sum = tf.reduce_mean(remote_b0)
  right_upper_mul = tf.multiply(right_upper_sum, casted_ids2)
  right_upper_add = tf.add(right_upper_mul, left_upper_sum)
  right_upper_round = tf.math.round(right_upper_mul)
  right_upper_floormod = tf.math.floormod(right_upper_round, N)

  left_upper_add = tf.add_n([left_upper_sum, casted_ids1, right_upper_add])
  left_upper_round = tf.math.round(left_upper_add)
  left_upper_floormod = tf.math.floormod(left_upper_round, N)

  remote_a1 = remote_op_a(left_upper_floormod)
  remote_b1 = remote_op_b(left_upper_floormod, right_upper_floormod)

  left_lower_sum = tf.reduce_mean(remote_a1)

  right_lower_sum = tf.reduce_mean(remote_b1)
  right_lower_mul = tf.multiply(casted_ids2, right_lower_sum)
  right_lower_div = tf.divide(right_upper_add, right_lower_mul)

  main_result = tf.add_n([
      left_lower_sum, right_lower_div, right_lower_sum, right_upper_sum,
      tf.cast(left_upper_floormod, tf.float32)
  ])
//  [python]
def generate_ignore_mask(labels,
                         ignore_labels):
  """Generates ignore mask used by contrastive loss.

  Args:
    labels: A tensor of shape [batch_size, height, width, 1]
    ignore_labels: A list of labels to ignore. Pixels with these labels will
      be ignored in the loss computation

  Returns:
    A tensor of shape [batch, num_pixels, num_pixels], indicating which pixel
      pairs are invalid.
  """
  ignore_mask = tf.math.reduce_any(
      tf.equal(labels, tf.constant(ignore_labels, dtype=labels.dtype)),
      axis=2, keepdims=True)
  ignore_mask = tf.to_float(
      tf.logical_or(ignore_mask, tf.transpose(ignore_mask, [0, 2, 1])))
  return ignore_mask
//  [python]
def generate_ignore_mask(labels,
                         ignore_labels):
  """Generates ignore mask used by contrastive loss.

  Args:
    labels: A tensor of shape [batch_size, height, width, 1]
    ignore_labels: A list of labels to ignore. Pixels with these labels will
      be ignored in the loss computation

  Returns:
    A tensor of shape [batch, num_pixels, num_pixels], indicating which pixel
      pairs are invalid.
  """
  ignore_mask = tf.math.reduce_any(
      tf.equal(labels, tf.constant(ignore_labels, dtype=labels.dtype)),
      axis=2, keepdims=True)
  ignore_mask = tf.cast(
      tf.logical_or(ignore_mask, tf.transpose(ignore_mask, [0, 2, 1])),
      tf.float32)
  return ignore_mask
//  [python]
def _create_model(self):
        TAG_COLUMN = ['tag_category_list', 'tag_brand_list']
        for key in TAG_COLUMN:
            self._feature[key] = tf.strings.split(self._feature[key], '|')

        key_dict = {}
        with self._make_scope('input_layer', self._bf16, self._input_layer_partitioner):
            print('Adaptive emb = ', self._adaptive_emb, 'TF = ', self.tf)
            if self._adaptive_emb and not self.tf:
                '''Adaptive Embedding Feature part 1 of 2'''
                print('Adaptive Embedding Feature part 1 of 2')
                adaptive_mask_tensors = {}
                for col in HASH_INPUTS:
                    adaptive_mask_tensors[col] = tf.ones([args.batch_size],
                                                          tf.int32)
                input_emb = tf.feature_column.input_layer(
                    self._feature,
                    self._feature_column,
                    adaptive_mask_tensors=adaptive_mask_tensors,
                    cols_to_output_tensors=key_dict)
            else:
                input_emb = tf.feature_column.input_layer(
                    self._feature,
                    self._feature_column,
                    cols_to_output_tensors=key_dict)

        with self._make_scope('PLE', self._bf16, self._dense_layer_partitioner):
            if self._bf16:
                input_emb = tf.cast(input_emb, dtype=tf.bfloat16)
            ple_inputs = [input_emb] * (self._num_tasks + 1)
            ple_outputs = []
            for i in range(self._num_layers):
                with tf.variable_scope(f'extraction_network_{i}'):
                    if i == self._num_layers - 1: # the last extraction net
                        ple_outputs = self.cgc_model(inputs=ple_inputs, level_name='level_'+str(i)+'_', is_last=True)
                    else:
                        ple_outputs = self.cgc_model(inputs=ple_inputs, level_name='level_'+str(i)+'_', is_last=False)
                        ple_inputs = ple_outputs

            towers=[]
            for i, tower in enumerate(self._towers):
                tower_name = tower[0]
                hidden_units = tower[2]

                with tf.variable_scope(tower_name, reuse=tf.AUTO_REUSE):
                    tower_output = self._dnn(ple_outputs[i], dnn_hidden_units=hidden_units, layer_name='tower_'+tower_name)
                    final_tower_predict = tf.layers.dense(inputs=tower_output,
                                                        units=1,
                                                        activation=None,
                                                        name=f'{tower_name}_output')
                    self._add_layer_summary(final_tower_predict, f'{tower_name}_output')
                    if self._bf16:
                        final_tower_predict = tf.cast(final_tower_predict, dtype=tf.float32)
                    towers.append(final_tower_predict)
            tower_stack = tf.stack(towers, axis=1)
            self._logits = tf.squeeze(tower_stack, [2])
            self.probability = tf.math.sigmoid(self._logits)
            self.output = tf.round(self.probability)
//  [python]
def scatter_flat(
    indices: Union[tf.Tensor, tf.Variable],
    updates: Union[tf.Tensor, tf.Variable],
    /,
    *,
    size: Optional[int] = None,
    reduction: str = "sum",
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    if indices.dtype != tf.int32 or indices.dtype != tf.int64:
        if indices.dtype in [tf.int8, tf.int16, tf.uint8, tf.uint16]:
            indices = tf.cast(indices, tf.int32)
        else:
            indices = tf.cast(indices, tf.int64)
    target = out
    target_given = ivy.exists(target)
    if ivy.exists(size) and ivy.exists(target):
        ivy.utils.assertions.check_equal(len(target.shape), 1, as_array=False)
        ivy.utils.assertions.check_equal(target.shape[0], size, as_array=False)
    if not target_given:
        target = tf.zeros([size], dtype=updates.dtype)
        res = tf.tensor_scatter_nd_update(target, tf.expand_dims(indices, -1), updates)
    elif reduction == "max":
        res = tf.tensor_scatter_nd_max(target, tf.expand_dims(indices, -1), updates)
    elif reduction == "min":
        res = tf.tensor_scatter_nd_min(target, tf.expand_dims(indices, -1), updates)
    elif reduction == "replace":
        res = tf.tensor_scatter_nd_update(target, tf.expand_dims(indices, -1), updates)
    elif reduction == "sum":
        res = tf.tensor_scatter_nd_add(target, tf.expand_dims(indices, -1), updates)
    else:
        raise ivy.utils.exceptions.IvyException(
            f'reduction is {reduction}, but it must be one of "sum", "min", "max" or'
            ' "replace"'
        )
    return ivy.inplace_update(out, res) if ivy.exists(out) else res
//  [python]
def generate_positive_and_negative_masks(
    labels):
  """Generates positive and negative masks used by contrastive loss.

  Args:
    labels: A tensor of shape [batch_size, height, width, 1]

  Returns:
    positive_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are positive pairs.
    negative_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are negative pairs.
  """
  positive_mask = tf.to_float(
      tf.equal(labels, tf.transpose(labels, [0, 2, 1])))
  negative_mask = 1 - positive_mask
  return positive_mask, negative_mask
//  [python]
def generate_positive_and_negative_masks(
    labels):
  """Generates positive and negative masks used by contrastive loss.

  Args:
    labels: A tensor of shape [batch_size, height, width, 1]

  Returns:
    positive_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are positive pairs.
    negative_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are negative pairs.
  """
  positive_mask = tf.cast(
      tf.equal(labels, tf.transpose(labels, [0, 2, 1])), tf.float32)
  negative_mask = 1 - positive_mask
  return positive_mask, negative_mask
//  [python]
def scatter_nd(
    indices: Union[tf.Tensor, tf.Variable],
    updates: Union[tf.Tensor, tf.Variable],
    /,
    shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None,
    *,
    reduction: str = "sum",
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    updates_dtype = updates.dtype
    if ivy.exists(out):
        dtype = ivy.promote_types(out.dtype, updates_dtype)
    updates = tf.cast(
        updates,
        (ivy.as_native_dtype(dtype) if ivy.exists(out) else updates_dtype),
    )

    expected_shape = (
        list(tf.shape(indices)[:-1]) + list(out.shape[tf.shape(indices)[-1] :])
        if ivy.exists(out)
        else list(tf.shape(indices)[:-1]) + list(shape[tf.shape(indices)[-1] :])
    )
    updates = _broadcast_to(updates, expected_shape)._data
    if len(updates.shape) == 0:
        indices = tf.expand_dims(indices, 0)
        updates = tf.expand_dims(updates, 0)

    # implementation
    target = out
    target_given = ivy.exists(target)
    if ivy.exists(shape) and target_given:
        ivy.utils.assertions.check_equal(
            ivy.Shape(target.shape), ivy.Shape(shape), as_array=False
        )
    if not target_given:
        shape = list(shape) if ivy.exists(shape) else list(out.shape)
        target = tf.zeros(shape, dtype=updates.dtype)
    if reduction == "sum":
        res = tf.tensor_scatter_nd_add(target, indices, updates)
    elif reduction == "min":
        res = tf.tensor_scatter_nd_min(target, indices, updates)
    elif reduction == "max":
        res = tf.tensor_scatter_nd_max(target, indices, updates)
    elif reduction == "mul":
        updates = ivy.multiply(ivy.gather_nd(target, indices), updates).data
        res = tf.tensor_scatter_nd_update(target, indices, updates)
    elif reduction == "replace":
        res = tf.tensor_scatter_nd_update(target, indices, updates)
    else:
        raise ivy.utils.exceptions.IvyException(
            f'reduction is {reduction}, but it must be one of "sum", "min", "max",'
            ' "mul" or "replace"'
        )
    if ivy.exists(out):
        return ivy.inplace_update(out, res)
    return res
//  [python]
def compute_contrastive_loss(logits,
                             positive_mask,
                             negative_mask,
                             ignore_mask):
  """Contrastive loss function.

  Args:
    logits: A tensor of shape [batch, num_pixels, num_pixels] with each value
      corresponding to a pixel pair.
    positive_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are positive pairs.
    negative_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are negative pairs.
    ignore_mask: A tensor of shape [batch, num_pixels, num_pixels], indicating
      which pixel pairs are invalid.

  Returns:
    A scalar tensor with contrastive loss
  """
  validity_mask = 1 - ignore_mask
  positive_mask *= validity_mask
  negative_mask *= validity_mask

  exp_logits = tf.exp(logits) * validity_mask

  normalized_exp_logits = tf.math.divide_no_nan(
      exp_logits,
      exp_logits + tf.reduce_sum(exp_logits * negative_mask, 2, keepdims=True))
  neg_log_likelihood = -tf.math.log(
      normalized_exp_logits * validity_mask + ignore_mask)

  # normalize weight and sum in dimension 2
  normalized_weight = positive_mask / tf.maximum(
      1e-6, tf.reduce_sum(positive_mask, 2, keepdims=True))
  neg_log_likelihood = tf.reduce_sum(
      neg_log_likelihood * normalized_weight, axis=2)

  # normalize weight and sum in dimension 1
  positive_mask_sum = tf.reduce_sum(positive_mask, 2)
  valid_index = 1 - tf.to_float(tf.equal(positive_mask_sum, 0))
  normalized_weight = valid_index / tf.maximum(
      1e-6, tf.reduce_sum(valid_index, 1, keepdims=True))
  neg_log_likelihood = tf.reduce_sum(
      neg_log_likelihood * normalized_weight, axis=1)
  loss = tf.reduce_mean(neg_log_likelihood)

  return loss
//  [python]
def compute_contrastive_loss(logits,
                             positive_mask,
                             negative_mask,
                             ignore_mask):
  """Contrastive loss function.

  Args:
    logits: A tensor of shape [batch, num_pixels, num_pixels] with each value
      corresponding to a pixel pair.
    positive_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are positive pairs.
    negative_mask: A tensor of shape [batch, num_pixels, num_pixels] indicating
      which pixel pairs are negative pairs.
    ignore_mask: A tensor of shape [batch, num_pixels, num_pixels], indicating
      which pixel pairs are invalid.

  Returns:
    A scalar tensor with contrastive loss
  """
  validity_mask = 1 - ignore_mask
  positive_mask *= validity_mask
  negative_mask *= validity_mask

  exp_logits = tf.exp(logits) * validity_mask

  normalized_exp_logits = tf.math.divide_no_nan(
      exp_logits,
      exp_logits + tf.reduce_sum(exp_logits * negative_mask, 2, keepdims=True))
  neg_log_likelihood = -tf.math.log(
      normalized_exp_logits * validity_mask + ignore_mask)

  # normalize weight and sum in dimension 2
  normalized_weight = positive_mask / tf.maximum(
      1e-6, tf.reduce_sum(positive_mask, 2, keepdims=True))
  neg_log_likelihood = tf.reduce_sum(
      neg_log_likelihood * normalized_weight, axis=2)

  # normalize weight and sum in dimension 1
  positive_mask_sum = tf.reduce_sum(positive_mask, 2)
  valid_index = 1 - tf.cast(tf.equal(positive_mask_sum, 0), tf.float32)
  normalized_weight = valid_index / tf.maximum(
      1e-6, tf.reduce_sum(valid_index, 1, keepdims=True))
  neg_log_likelihood = tf.reduce_sum(
      neg_log_likelihood * normalized_weight, axis=1)
  loss = tf.reduce_mean(neg_log_likelihood)

  return loss
//  [python]
def predict_proba_incremental(self, waveform,upsampled_local_condition=None, global_condition=None,name='wavenet'):
        """
        local_condition: upsampled local condition
        """


        with tf.variable_scope(name):
            
            if self.scalar_input:
                encoded = tf.reshape(waveform , [self.batch_size, -1, 1])
            else:
                encoded = tf.one_hot(waveform, self.quantization_channels)
                encoded = tf.reshape(encoded, [self.batch_size,-1, self.quantization_channels])   # encoded shape=(N,1, 256)
            
            gc_embedding = self._embed_gc(global_condition)                   # --> shape=(1, 1, 32)
            
            
            # local condition
            if upsampled_local_condition is not None:
                upsampled_local_condition = tf.reshape(upsampled_local_condition , [self.batch_size, -1, self.local_condition_channels])
            
            raw_output = self._create_network(encoded,upsampled_local_condition,gc_embedding)        # 이것이 fast generation algorithm의 핵심  --> (batch_size, 1, 256)
            
            if self.scalar_input:
                out = tf.reshape(raw_output, [self.batch_size, -1, self.out_channels])
                proba = sample_from_discretized_mix_logistic(out)
            else:
                out = tf.reshape(raw_output, [self.batch_size, self.quantization_channels])
                proba = tf.cast(tf.nn.softmax(tf.cast(out, tf.float64)), tf.float32)

            return proba
//  [python]
def convert_ms_to_frame_number(batch, spectrogram_feature_name):
  """Convert start timestamps in microsec into the frame numbers."""
  word_start_timestamp = batch['mwp_start_timestamp']
  word_end_timestamp = batch['mwp_end_timestamp']
  wave_length = batch['mwp_wave_length']
  spectrogram = batch[spectrogram_feature_name]

  start_idxs = t.cast(
      t.math.divide_no_nan(word_start_timestamp, wave_length)
      * t.cast((t.shape(spectrogram)[0] - 1), t.float32), t.int32)
  end_idxs = t.cast(
      t.math.divide_no_nan(word_end_timestamp, wave_length)
      * t.cast((t.shape(spectrogram)[0] - 1), t.float32), t.int32)

  batch['mwp_start_indices'] = start_idxs
  batch['mwp_end_indices'] = end_idxs

  del batch['mwp_wave_length'], batch['mwp_start_timestamp'], batch[
      'mwp_end_timestamp']

  return batch
//  [python]
def convert_ms_to_frame_number(batch, spectrogram_feature_name):
  """Convert start timestamps in microsec into the frame numbers."""
  word_start_timestamp = batch['mwp_start_timestamp']
  word_end_timestamp = batch['mwp_end_timestamp']
  wave_length = batch['mwp_wave_length']
  spectrogram = batch[spectrogram_feature_name]

  start_idxs = tf.cast(
      tf.math.divide_no_nan(word_start_timestamp, wave_length)
      * tf.cast((tf.shape(spectrogram)[0] - 1), tf.float32), tf.int32)
  end_idxs = tf.cast(
      tf.math.divide_no_nan(word_end_timestamp, wave_length)
      * tf.cast((tf.shape(spectrogram)[0] - 1), tf.float32), tf.int32)

  batch['mwp_start_indices'] = start_idxs
  batch['mwp_end_indices'] = end_idxs

  del batch['mwp_wave_length'], batch['mwp_start_timestamp'], batch[
      'mwp_end_timestamp']

  return batch
//  [python]
def add_white_noise(wave_tensor, snr):
  """Adds deterministic gaussian white noise to the signal.

  While setting a global seeds does not guarrantee the deterministic per-sample
  noise generation due to the intrinsic randomness in the TF dataloading
  process. Instead, we generate the seeds from the input signal and generate
  noise from the stateless normal sampling process. This effectively sample
  deterministic noise per sample while the noise varies between samples. (A
  constant seed for all sample would generate identical noise for all samples,
  which is not a desired behavior.)

  Args:
    wave_tensor: Input tensor containing the raw wave signal.
    snr: The signal to noise ratio to determine the power of the noise.

  Returns:
    Output signal with gaussian white noise injected.
  """
  pos = tf.reduce_mean(tf.square(wave_tensor))
  nos = pos / snr
  # Deterministically create the noise using the input signal as the seeds.
  # The seeds are determined by the sums of the positive and negative signals.
  positive_sum = t.cast(tf.reduce_sum(tf.maximum(wave_tensor, 0)), t.int32)
  negative_sum = t.cast(tf.reduce_sum(tf.maximum(-wave_tensor, 0)), t.int32)
  noise = tf.random.stateless_normal(
      tf.shape(wave_tensor),
      seed=[positive_sum, negative_sum],
      stddev=tf.sqrt(nos),
      dtype=wave_tensor.dtype)

  return wave_tensor + noise
//  [python]
def add_white_noise(wave_tensor, snr):
  """Adds deterministic gaussian white noise to the signal.

  While setting a global seeds does not guarrantee the deterministic per-sample
  noise generation due to the intrinsic randomness in the TF dataloading
  process. Instead, we generate the seeds from the input signal and generate
  noise from the stateless normal sampling process. This effectively sample
  deterministic noise per sample while the noise varies between samples. (A
  constant seed for all sample would generate identical noise for all samples,
  which is not a desired behavior.)

  Args:
    wave_tensor: Input tensor containing the raw wave signal.
    snr: The signal to noise ratio to determine the power of the noise.

  Returns:
    Output signal with gaussian white noise injected.
  """
  pos = tf.reduce_mean(tf.square(wave_tensor))
  nos = pos / snr
  # Deterministically create the noise using the input signal as the seeds.
  # The seeds are determined by the sums of the positive and negative signals.
  positive_sum = tf.cast(tf.reduce_sum(tf.maximum(wave_tensor, 0)), tf.int32)
  negative_sum = tf.cast(tf.reduce_sum(tf.maximum(-wave_tensor, 0)), tf.int32)
  noise = tf.random.stateless_normal(
      tf.shape(wave_tensor),
      seed=[positive_sum, negative_sum],
      stddev=tf.sqrt(nos),
      dtype=wave_tensor.dtype)

  return wave_tensor + noise