jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _head(self, policy_input, heading, xy, target_xy):
    """Build the head of the agent: linear policy and value function, and pass
    the auxiliary outputs through.
    """

    # Linear policy and value function.
    policy_logits = snt.Linear(
        self._num_actions, name='policy_logits')(policy_input)
    baseline = tf.squeeze(snt.Linear(1, name='baseline')(policy_input), axis=-1)

    # Sample an action from the policy.
    new_action = tf.random.categorical(
        policy_logits, num_samples=1, dtype=tf.int32)
    new_action = tf.squeeze(new_action, 1, name='new_action')

    return AgentOutput(
        new_action, policy_logits, baseline, heading, xy, target_xy)
//  [python]
def symbolic_translate(self, inp, greedy=False, max_len = None, eps = 1e-30, **flags):
        """
        takes symbolic int32 matrix of hebrew words, produces output tokens sampled
        from the model and output log-probabilities for all possible tokens at each tick.
        :param inp: input sequence, int32 matrix of shape [batch,time]
        :param greedy: if greedy, takes token with highest probablity at each tick.
            Otherwise samples proportionally to probability.
        :param max_len: max length of output, defaults to 2 * input length
        :return: output tokens int32[batch,time] and
                 log-probabilities of all tokens at each tick, [batch,time,n_tokens]
        """
        first_state = self.encode(inp, **flags)

        batch_size = tf.shape(inp)[0]
        bos = tf.fill([batch_size],self.out_voc.bos_ix)
        first_logits = tf.log(tf.one_hot(bos, len(self.out_voc)) + eps)
        max_len = tf.reduce_max(tf.shape(inp)[1])*2

        def step(blob,t):
            h_prev, y_prev = blob[:-2], blob[-1]
            h_new, logits = self.decode(h_prev, y_prev, **flags)
            y_new = tf.argmax(logits,axis=-1) if greedy else tf.multinomial(logits,1)[:,0]
            return list(h_new) + [logits, tf.cast(y_new,y_prev.dtype)]

        results = tf.scan(step, initializer=list(first_state) + [first_logits, bos],
                          elems=[tf.range(max_len)])

        # gather state, logits and outs, each of shape [time,batch,...]
        states_seq, logits_seq, out_seq = results[:-2], results[-2], results[-1]

        # add initial state, logits and out
        logits_seq = tf.concat((first_logits[None],logits_seq),axis=0)
        out_seq = tf.concat((bos[None], out_seq), axis=0)
        states_seq = [tf.concat((init[None], states), axis=0)
                      for init, states in zip(first_state, states_seq)]

        #convert from [time,batch,...] to [batch,time,...]
        logits_seq = tf.transpose(logits_seq, [1, 0, 2])
        out_seq = tf.transpose(out_seq)
        states_seq = [tf.transpose(states, [1, 0] + list(range(2, states.shape.ndims)))
                      for states in states_seq]

        return out_seq, tf.nn.log_softmax(logits_seq)
//  [python]
def symbolic_translate(self, inp, greedy=False, max_len = None, eps = 1e-30, **flags):
        """
        takes symbolic int32 matrix of hebrew words, produces output tokens sampled
        from the model and output log-probabilities for all possible tokens at each tick.
        :param inp: input sequence, int32 matrix of shape [batch,time]
        :param greedy: if greedy, takes token with highest probablity at each tick.
            Otherwise samples proportionally to probability.
        :param max_len: max length of output, defaults to 2 * input length
        :return: output tokens int32[batch,time] and
                 log-probabilities of all tokens at each tick, [batch,time,n_tokens]
        """
        first_state = self.encode(inp, **flags)

        batch_size = tf.shape(inp)[0]
        bos = tf.fill([batch_size],self.out_voc.bos_ix)
        first_logits = tf.log(tf.one_hot(bos, len(self.out_voc)) + eps)
        max_len = tf.reduce_max(tf.shape(inp)[1])*2

        def step(blob,t):
            h_prev, y_prev = blob[:-2], blob[-1]
            h_new, logits = self.decode(h_prev, y_prev, **flags)
            y_new = tf.argmax(logits,axis=-1) if greedy else tf.random.categorical(logits,1)[:,0]
            return list(h_new) + [logits, tf.cast(y_new,y_prev.dtype)]

        results = tf.scan(step, initializer=list(first_state) + [first_logits, bos],
                          elems=[tf.range(max_len)])

        # gather state, logits and outs, each of shape [time,batch,...]
        states_seq, logits_seq, out_seq = results[:-2], results[-2], results[-1]

        # add initial state, logits and out
        logits_seq = tf.concat((first_logits[None],logits_seq),axis=0)
        out_seq = tf.concat((bos[None], out_seq), axis=0)
        states_seq = [tf.concat((init[None], states), axis=0)
                      for init, states in zip(first_state, states_seq)]

        #convert from [time,batch,...] to [batch,time,...]
        logits_seq = tf.transpose(logits_seq, [1, 0, 2])
        out_seq = tf.transpose(out_seq)
        states_seq = [tf.transpose(states, [1, 0] + list(range(2, states.shape.ndims)))
                      for states in states_seq]

        return out_seq, tf.nn.log_softmax(logits_seq)
//  [python]
def mutate_x(x, mutation_rate, seq_len):
  """Randomly and independently mutate a sequence based on a mutation rate."""
  # generate mutations for all positions,
  # in order to be different than itselves, the mutations have to be >= 1
  # mute the untargeted positions by multiple mask (1 for targeted)
  # adding the mutations to the original, mod 4 if necessary
  tf.random.set_random_seed(time.time())

  def true_fn():
    """no mutations."""
    return x

  def false_fn():
    """add mutations."""
    mask = tf.cast(
        tf.multinomial(tf.log([[1 - mutation_rate, mutation_rate]]), seq_len),
        tf.int32)[0]
    possible_mutations = tf.random_uniform([seq_len],
                                           minval=1,
                                           maxval=4,
                                           dtype=tf.int32)
    x_new = tf.mod(x + mask * possible_mutations, 4)
    return x_new

  return tf.cond(tf.equal(mutation_rate, 0), true_fn, false_fn)
//  [python]
def mutate_x(x, mutation_rate, seq_len):
  """Randomly and independently mutate a sequence based on a mutation rate."""
  # generate mutations for all positions,
  # in order to be different than itselves, the mutations have to be >= 1
  # mute the untargeted positions by multiple mask (1 for targeted)
  # adding the mutations to the original, mod 4 if necessary
  tf.random.set_random_seed(time.time())

  def true_fn():
    """no mutations."""
    return x

  def false_fn():
    """add mutations."""
    mask = tf.cast(
        tf.random.categorical(tf.log([[1 - mutation_rate, mutation_rate]]), seq_len),
        tf.int32)[0]
    possible_mutations = tf.random.uniform([seq_len],
                                           minval=1,
                                           maxval=4,
                                           dtype=tf.int32)
    x_new = tf.mod(x + mask * possible_mutations, 4)
    return x_new

  return tf.cond(tf.equal(mutation_rate, 0), true_fn, false_fn)
//  [python]
def _legal_sample(probs, legal_action=None, use_max=False):
    # Note: just sample for now
    if use_max:
        return tf.reshape(tf.argmax(probs, axis=1), [-1])
    return tf.reshape(tf.multinomial(tf.log(probs), 1), [-1])
//  [python]
def _legal_sample(probs, legal_action=None, use_max=False):
    # Note: just sample for now
    if use_max:
        return tf.reshape(tf.argmax(probs, axis=1), [-1])
    return tf.reshape(tf.random.categorical(tf.log(probs), 1), [-1])
//  [python]
def _create_network(self, feature_space, new_space, gather_space):
        # input
        input_feature = tf.placeholder(tf.float32, (None,) + feature_space)
        input_feature2 = tf.placeholder(tf.float32, (None,) + new_space)
        input_feature3 = tf.placeholder(tf.float32, (None,) + gather_space)
        input_act_prob = tf.placeholder(tf.float32, (None, self.num_actions))
        action = tf.placeholder(tf.int32, [None])

        reward = tf.placeholder(tf.float32, [None])

        hidden_size = [256]

        # fully connected

        h_emb = tf.layers.dense(input_feature,  units=hidden_size[0], activation=tf.nn.relu)
        h_emb2 = tf.layers.dense(input_feature2,  units=hidden_size[0], activation=tf.nn.relu)

        concat_layer = tf.concat([h_emb, h_emb2], axis=1)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0] * 2, activation=tf.nn.relu)

        policy = tf.layers.dense(dense / 0.1, units=self.num_actions, activation=tf.nn.softmax)
        policy = tf.clip_by_value(policy, 1e-10, 1-1e-10)

        self.calc_action = tf.multinomial(tf.log(policy), 1)

        # for value obtain
        emb_prob = tf.layers.dense(input_act_prob, units=64, activation=tf.nn.relu)
        dense_prob = tf.layers.dense(emb_prob, units=32, activation=tf.nn.relu)
        concat_layer = tf.concat([concat_layer, dense_prob], axis=1)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0], activation=tf.nn.relu)
        value = tf.layers.dense(dense, units=1)
        value = tf.reshape(value, (-1,))

        action_mask = tf.one_hot(action, self.num_actions)
        advantage = tf.stop_gradient(reward - value)

        log_policy = tf.log(policy + 1e-6)
        log_prob = tf.reduce_sum(log_policy * action_mask, axis=1)

        pg_loss = -tf.reduce_mean(advantage * log_prob)
        vf_loss = self.value_coef * tf.reduce_mean(tf.square(reward - value))
        neg_entropy = self.ent_coef * tf.reduce_mean(tf.reduce_sum(policy * log_policy, axis=1))
        total_loss = pg_loss + vf_loss + neg_entropy

        # train op (clip gradient)
        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
        gradients, variables = zip(*optimizer.compute_gradients(total_loss))
        gradients, _ = tf.clip_by_global_norm(gradients, 5.0)
        self.train_op = optimizer.apply_gradients(zip(gradients, variables))

        train_op = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(total_loss)

        self.input_feature = input_feature
        self.input_feature2 = input_feature2
        self.input_feature3 = input_feature3
        self.input_act_prob = input_act_prob
        self.action = action
        self.reward = reward

        self.policy, self.value = policy, value
        self.train_op = train_op
        self.pg_loss, self.vf_loss, self.reg_loss = pg_loss, vf_loss, neg_entropy
        self.total_loss = total_loss
//  [python]
def _create_network(self, feature_space, new_space, gather_space):
        # input
        input_feature = tf.placeholder(tf.float32, (None,) + feature_space)
        input_feature2 = tf.placeholder(tf.float32, (None,) + new_space)
        input_feature3 = tf.placeholder(tf.float32, (None,) + gather_space)
        input_act_prob = tf.placeholder(tf.float32, (None, self.num_actions))
        action = tf.placeholder(tf.int32, [None])

        reward = tf.placeholder(tf.float32, [None])

        hidden_size = [256]

        # fully connected

        h_emb = tf.layers.dense(input_feature,  units=hidden_size[0], activation=tf.nn.relu)
        h_emb2 = tf.layers.dense(input_feature2,  units=hidden_size[0], activation=tf.nn.relu)

        concat_layer = tf.concat([h_emb, h_emb2], axis=1)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0] * 2, activation=tf.nn.relu)

        policy = tf.layers.dense(dense / 0.1, units=self.num_actions, activation=tf.nn.softmax)
        policy = tf.clip_by_value(policy, 1e-10, 1-1e-10)

        self.calc_action = tf.random.categorical(tf.log(policy), 1)

        # for value obtain
        emb_prob = tf.layers.dense(input_act_prob, units=64, activation=tf.nn.relu)
        dense_prob = tf.layers.dense(emb_prob, units=32, activation=tf.nn.relu)
        concat_layer = tf.concat([concat_layer, dense_prob], axis=1)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0], activation=tf.nn.relu)
        value = tf.layers.dense(dense, units=1)
        value = tf.reshape(value, (-1,))

        action_mask = tf.one_hot(action, self.num_actions)
        advantage = tf.stop_gradient(reward - value)

        log_policy = tf.log(policy + 1e-6)
        log_prob = tf.reduce_sum(log_policy * action_mask, axis=1)

        pg_loss = -tf.reduce_mean(advantage * log_prob)
        vf_loss = self.value_coef * tf.reduce_mean(tf.square(reward - value))
        neg_entropy = self.ent_coef * tf.reduce_mean(tf.reduce_sum(policy * log_policy, axis=1))
        total_loss = pg_loss + vf_loss + neg_entropy

        # train op (clip gradient)
        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
        gradients, variables = zip(*optimizer.compute_gradients(total_loss))
        gradients, _ = tf.clip_by_global_norm(gradients, 5.0)
        self.train_op = optimizer.apply_gradients(zip(gradients, variables))

        train_op = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(total_loss)

        self.input_feature = input_feature
        self.input_feature2 = input_feature2
        self.input_feature3 = input_feature3
        self.input_act_prob = input_act_prob
        self.action = action
        self.reward = reward

        self.policy, self.value = policy, value
        self.train_op = train_op
        self.pg_loss, self.vf_loss, self.reg_loss = pg_loss, vf_loss, neg_entropy
        self.total_loss = total_loss
//  [python]
def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0, top_p=0.0):
    if start_token is None:
        assert context is not None, 'Specify exactly one of start_token and context!'
    else:
        assert context is None, 'Specify exactly one of start_token and context!'
        context = tf.fill([batch_size, 1], start_token)

    def step(hparams, tokens, past=None):
        lm_output = model.model(hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)

        logits = lm_output['logits'][:, :, :hparams.n_vocab]
        presents = lm_output['present']
        presents.set_shape(model.past_shape(hparams=hparams, batch_size=batch_size))
        return {
            'logits': logits,
            'presents': presents,
        }

    with tf.name_scope('sample_sequence'):
        def body(past, prev, output):
            next_outputs = step(hparams, prev, past=past)
            logits = next_outputs['logits'][:, -1, :]  / tf.to_float(temperature)
            if top_p > 0.0:
                logits = top_p_logits(logits, p=top_p)
            else:
                logits = top_k_logits(logits, k=top_k)
            samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32)
            return [
                next_outputs['presents'] if past is None else tf.concat([past, next_outputs['presents']], axis=-2),
                samples,
                tf.concat([output, samples], axis=1)
            ]

        past, prev, output = body(None, context, context)

        def cond(*args):
            return True

        _, _, tokens = tf.while_loop(
            cond=cond, body=body,
            maximum_iterations=length - 1,
            loop_vars=[
                past,
                prev,
                output
            ],
            shape_invariants=[
                tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
                tf.TensorShape([batch_size, None]),
                tf.TensorShape([batch_size, None]),
            ],
            back_prop=False,
        )

        return tokens
//  [python]
def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0, top_p=0.0):
    if start_token is None:
        assert context is not None, 'Specify exactly one of start_token and context!'
    else:
        assert context is None, 'Specify exactly one of start_token and context!'
        context = tf.fill([batch_size, 1], start_token)

    def step(hparams, tokens, past=None):
        lm_output = model.model(hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)

        logits = lm_output['logits'][:, :, :hparams.n_vocab]
        presents = lm_output['present']
        presents.set_shape(model.past_shape(hparams=hparams, batch_size=batch_size))
        return {
            'logits': logits,
            'presents': presents,
        }

    with tf.name_scope('sample_sequence'):
        def body(past, prev, output):
            next_outputs = step(hparams, prev, past=past)
            logits = next_outputs['logits'][:, -1, :]  / tf.to_float(temperature)
            if top_p > 0.0:
                logits = top_p_logits(logits, p=top_p)
            else:
                logits = top_k_logits(logits, k=top_k)
            samples = tf.random.categorical(logits, num_samples=1)
            return [
                next_outputs['presents'] if past is None else tf.concat([past, next_outputs['presents']], axis=-2),
                samples,
                tf.concat([output, samples], axis=1)
            ]

        past, prev, output = body(None, context, context)

        def cond(*args):
            return True

        _, _, tokens = tf.while_loop(
            cond=cond, body=body,
            maximum_iterations=length - 1,
            loop_vars=[
                past,
                prev,
                output
            ],
            shape_invariants=[
                tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
                tf.TensorShape([batch_size, None]),
                tf.TensorShape([batch_size, None]),
            ],
            back_prop=False,
        )

        return tokens
//  [python]
def multinomial_squeeze(self, logits, temperature = 1.0):
        """multinomial sampling from logits."""
        logits_shape = utils.shape_list(logits)
        reshaped_logits = (
            tf.reshape(logits, [-1, logits_shape[-1]]) / temperature
        )
        choices = tf.multinomial(reshaped_logits, 1)
        choices = tf.reshape(choices, logits_shape[:-1])
        return tf.to_int32(choices)
//  [python]
def multinomial_squeeze(self, logits, temperature = 1.0):
        """multinomial sampling from logits."""
        logits_shape = utils.shape_list(logits)
        reshaped_logits = (
            tf.reshape(logits, [-1, logits_shape[-1]]) / temperature
        )
        choices = tf.random.categorical(reshaped_logits, 1)
        choices = tf.reshape(choices, logits_shape[:-1])
        return tf.to_int32(choices)
//  [python]
def _create_network(self, feature_space, new_space, gather_space):
        # input
        input_feature = tf.placeholder(tf.float32, (None,) + feature_space)
        input_feature2 = tf.placeholder(tf.float32, (None,) + new_space)
        input_feature3 = tf.placeholder(tf.float32, (None,) + gather_space)
        input_act_prob = tf.placeholder(tf.float32, (None, self.num_actions))
        action = tf.placeholder(tf.int32, [None])

        reward = tf.placeholder(tf.float32, [None])

        hidden_size = [256]

        # fully connected

        h_emb = tf.layers.dense(input_feature,  units=hidden_size[0], activation=tf.nn.relu)
        h_emb2 = tf.layers.dense(input_feature2,  units=hidden_size[0], activation=tf.nn.relu)

        concat_layer = tf.concat([h_emb, h_emb2], axis=1)

        h_emb2 = tf.layers.dense(input_feature2,  units=hidden_size[0], activation=tf.nn.relu)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0] * 2, activation=tf.nn.relu)

        policy = tf.layers.dense(dense / 0.1, units=self.num_actions, activation=tf.nn.softmax)
        policy = tf.clip_by_value(policy, 1e-10, 1-1e-10)

        self.calc_action = tf.multinomial(tf.log(policy), 1)

        # for value obtain
        emb_prob = tf.layers.dense(input_act_prob, units=64, activation=tf.nn.relu)
        dense_prob = tf.layers.dense(emb_prob, units=32, activation=tf.nn.relu)
        concat_layer = tf.concat([concat_layer, dense_prob], axis=1)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0], activation=tf.nn.relu)
        value = tf.layers.dense(dense, units=1)
        value = tf.reshape(value, (-1,))

        action_mask = tf.one_hot(action, self.num_actions)
        advantage = tf.stop_gradient(reward - value)

        log_policy = tf.log(policy + 1e-6)
        log_prob = tf.reduce_sum(log_policy * action_mask, axis=1)

        pg_loss = -tf.reduce_mean(advantage * log_prob)
        vf_loss = self.value_coef * tf.reduce_mean(tf.square(reward - value))
        neg_entropy = self.ent_coef * tf.reduce_mean(tf.reduce_sum(policy * log_policy, axis=1))
        total_loss = pg_loss + vf_loss + neg_entropy

        # train op (clip gradient)
        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
        gradients, variables = zip(*optimizer.compute_gradients(total_loss))
        gradients, _ = tf.clip_by_global_norm(gradients, 5.0)
        self.train_op = optimizer.apply_gradients(zip(gradients, variables))

        train_op = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(total_loss)

        self.input_feature = input_feature
        self.input_feature2 = input_feature2
        self.input_feature3 = input_feature3
        self.input_act_prob = input_act_prob
        self.action = action
        self.reward = reward

        self.policy, self.value = policy, value
        self.train_op = train_op
        self.pg_loss, self.vf_loss, self.reg_loss = pg_loss, vf_loss, neg_entropy
        self.total_loss = total_loss
//  [python]
def _create_network(self, feature_space, new_space, gather_space):
        # input
        input_feature = tf.placeholder(tf.float32, (None,) + feature_space)
        input_feature2 = tf.placeholder(tf.float32, (None,) + new_space)
        input_feature3 = tf.placeholder(tf.float32, (None,) + gather_space)
        input_act_prob = tf.placeholder(tf.float32, (None, self.num_actions))
        action = tf.placeholder(tf.int32, [None])

        reward = tf.placeholder(tf.float32, [None])

        hidden_size = [256]

        # fully connected

        h_emb = tf.layers.dense(input_feature,  units=hidden_size[0], activation=tf.nn.relu)
        h_emb2 = tf.layers.dense(input_feature2,  units=hidden_size[0], activation=tf.nn.relu)

        concat_layer = tf.concat([h_emb, h_emb2], axis=1)

        h_emb2 = tf.layers.dense(input_feature2,  units=hidden_size[0], activation=tf.nn.relu)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0] * 2, activation=tf.nn.relu)

        policy = tf.layers.dense(dense / 0.1, units=self.num_actions, activation=tf.nn.softmax)
        policy = tf.clip_by_value(policy, 1e-10, 1-1e-10)

        self.calc_action = tf.random.categorical(tf.log(policy), 1)

        # for value obtain
        emb_prob = tf.layers.dense(input_act_prob, units=64, activation=tf.nn.relu)
        dense_prob = tf.layers.dense(emb_prob, units=32, activation=tf.nn.relu)
        concat_layer = tf.concat([concat_layer, dense_prob], axis=1)
        dense = tf.layers.dense(concat_layer, units=hidden_size[0], activation=tf.nn.relu)
        value = tf.layers.dense(dense, units=1)
        value = tf.reshape(value, (-1,))

        action_mask = tf.one_hot(action, self.num_actions)
        advantage = tf.stop_gradient(reward - value)

        log_policy = tf.log(policy + 1e-6)
        log_prob = tf.reduce_sum(log_policy * action_mask, axis=1)

        pg_loss = -tf.reduce_mean(advantage * log_prob)
        vf_loss = self.value_coef * tf.reduce_mean(tf.square(reward - value))
        neg_entropy = self.ent_coef * tf.reduce_mean(tf.reduce_sum(policy * log_policy, axis=1))
        total_loss = pg_loss + vf_loss + neg_entropy

        # train op (clip gradient)
        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
        gradients, variables = zip(*optimizer.compute_gradients(total_loss))
        gradients, _ = tf.clip_by_global_norm(gradients, 5.0)
        self.train_op = optimizer.apply_gradients(zip(gradients, variables))

        train_op = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(total_loss)

        self.input_feature = input_feature
        self.input_feature2 = input_feature2
        self.input_feature3 = input_feature3
        self.input_act_prob = input_act_prob
        self.action = action
        self.reward = reward

        self.policy, self.value = policy, value
        self.train_op = train_op
        self.pg_loss, self.vf_loss, self.reg_loss = pg_loss, vf_loss, neg_entropy
        self.total_loss = total_loss
//  [python]
def provide_one_hot_labels(self, batch_size):
    """Provides one hot labels."""
    pitch_counts = self.get_pitch_counts()
    pitches = sorted(pitch_counts.keys())
    counts = [pitch_counts[p] for p in pitches]
    indices = tf.reshape(
        tf.multinomial(tf.log([tf.to_float(counts)]), batch_size), [batch_size])
    one_hot_labels = tf.one_hot(indices, depth=len(pitches))
    return one_hot_labels
//  [python]
def provide_one_hot_labels(self, batch_size):
    """Provides one hot labels."""
    pitch_counts = self.get_pitch_counts()
    pitches = sorted(pitch_counts.keys())
    counts = [pitch_counts[p] for p in pitches]
    indices = tf.reshape(
        tf.random.categorical(tf.math.log([tf.cast(counts, tf.float32)]), batch_size), [batch_size])
    one_hot_labels = tf.one_hot(indices, depth=len(pitches))
    return one_hot_labels
//  [python]
def _head(self, core_output):
    """Build the head of the agent: linear policy and value function."""
    policy_logits = snt.Linear(
        self._num_actions, name='policy_logits')(
            core_output)
    baseline = tf.squeeze(snt.Linear(1, name='baseline')(core_output), axis=-1)

    # Sample an action from the policy.
    new_action = tf.multinomial(
        policy_logits, num_samples=1, output_dtype=tf.int32)
    new_action = tf.squeeze(new_action, 1, name='new_action')

    return AgentOutput(new_action, policy_logits, baseline)
//  [python]
def _head(self, core_output):
    """Build the head of the agent: linear policy and value function."""
    policy_logits = snt.Linear(
        self._num_actions, name='policy_logits')(
            core_output)
    baseline = tf.squeeze(snt.Linear(1, name='baseline')(core_output), axis=-1)

    # Sample an action from the policy.
    new_action = tf.random.categorical(
        policy_logits, num_samples=1, dtype=tf.int32)
    new_action = tf.squeeze(new_action, 1, name='new_action')

    return AgentOutput(new_action, policy_logits, baseline)
//  [python]
def vq_nearest_neighbor(x, means,
                        soft_em=False, num_samples=10, temperature=None):
  """Find the nearest element in means to elements in x."""
  bottleneck_size = common_layers.shape_list(means)[0]
  x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
  means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
  scalar_prod = tf.matmul(x, means, transpose_b=True)
  dist = x_norm_sq + tf.transpose(means_norm_sq) - 2 * scalar_prod
  if soft_em:
    x_means_idx = tf.multinomial(-dist, num_samples=num_samples)
    x_means_hot = tf.one_hot(
        x_means_idx, depth=common_layers.shape_list(means)[0])
    x_means_hot = tf.reduce_mean(x_means_hot, axis=1)
  else:
    if temperature is None:
      x_means_idx = tf.argmax(-dist, axis=-1)
    else:
      x_means_idx = tf.multinomial(- dist / temperature, 1)
      x_means_idx = tf.squeeze(x_means_idx, axis=-1)
    if (common_layers.should_generate_summaries() and
        not common_layers.is_xla_compiled()):
      tf.summary.histogram("means_idx", tf.reshape(x_means_idx, [-1]))
    x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
  x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
  x_means = tf.matmul(x_means_hot_flat, means)
  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
  return x_means_hot, e_loss, dist
//  [python]
def vq_nearest_neighbor(x, means,
                        soft_em=False, num_samples=10, temperature=None):
  """Find the nearest element in means to elements in x."""
  bottleneck_size = common_layers.shape_list(means)[0]
  x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
  means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
  scalar_prod = tf.matmul(x, means, transpose_b=True)
  dist = x_norm_sq + tf.transpose(means_norm_sq) - 2 * scalar_prod
  if soft_em:
    x_means_idx = tf.random.categorical(-dist, num_samples=num_samples)
    x_means_hot = tf.one_hot(
        x_means_idx, depth=common_layers.shape_list(means)[0])
    x_means_hot = tf.reduce_mean(x_means_hot, axis=1)
  else:
    if temperature is None:
      x_means_idx = tf.argmax(-dist, axis=-1)
    else:
      x_means_idx = tf.random.categorical(- dist / temperature, 1)
      x_means_idx = tf.squeeze(x_means_idx, axis=-1)
    if (common_layers.should_generate_summaries() and
        not common_layers.is_xla_compiled()):
      tf.summary.histogram("means_idx", tf.reshape(x_means_idx, [-1]))
    x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
  x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
  x_means = tf.matmul(x_means_hot_flat, means)
  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
  return x_means_hot, e_loss, dist
//  [python]
def sample(self, features):
        """Run the model and extract samples.

    Args:
      features: an map of string to `Tensor`.

    Returns:
       samples: an integer `Tensor`.
       logits: a list of `Tensor`s, one per datashard.
       losses: a dictionary: {loss-name (string): floating point `Scalar`}.
    """
        logits, losses = self(features)  # pylint: disable=not-callable
        if self._target_modality_is_real:
            return (
                logits,
                logits,
                losses,
            )  # Raw numbers returned from real modality.
        if self.hparams.sampling_method == 'argmax':
            samples = tf.argmax(logits, axis = -1)
        else:
            assert self.hparams.sampling_method == 'random'

            def multinomial_squeeze(logits, temperature = 1.0):
                logits_shape = common_layers.shape_list(logits)
                logits /= tf.reshape(
                    temperature, [-1] + [1] * (len(logits_shape) - 1)
                )
                reshaped_logits = tf.reshape(logits, [-1, logits_shape[-1]])
                choices = tf.multinomial(reshaped_logits, 1)
                choices = tf.reshape(choices, logits_shape[:-1])
                return choices

            temperature = features.get(
                'sampling_temp', self.hparams.sampling_temp
            )
            samples = multinomial_squeeze(logits, temperature)

        return samples, logits, losses
//  [python]
def sample(self, features):
        """Run the model and extract samples.

    Args:
      features: an map of string to `Tensor`.

    Returns:
       samples: an integer `Tensor`.
       logits: a list of `Tensor`s, one per datashard.
       losses: a dictionary: {loss-name (string): floating point `Scalar`}.
    """
        logits, losses = self(features)  # pylint: disable=not-callable
        if self._target_modality_is_real:
            return (
                logits,
                logits,
                losses,
            )  # Raw numbers returned from real modality.
        if self.hparams.sampling_method == 'argmax':
            samples = tf.argmax(logits, axis = -1)
        else:
            assert self.hparams.sampling_method == 'random'

            def multinomial_squeeze(logits, temperature = 1.0):
                logits_shape = common_layers.shape_list(logits)
                logits /= tf.reshape(
                    temperature, [-1] + [1] * (len(logits_shape) - 1)
                )
                reshaped_logits = tf.reshape(logits, [-1, logits_shape[-1]])
                choices = tf.random.categorical(reshaped_logits, 1)
                choices = tf.reshape(choices, logits_shape[:-1])
                return choices

            temperature = features.get(
                'sampling_temp', self.hparams.sampling_temp
            )
            samples = multinomial_squeeze(logits, temperature)

        return samples, logits, losses
//  [python]
def sample(self, features):
    """Run the model and extract samples.

    Args:
      features: an map of string to `Tensor`.

    Returns:
       samples: an integer `Tensor`.
       logits: a list of `Tensor`s, one per datashard.
       losses: a dictionary: {loss-name (string): floating point `Scalar`}.
    """
    logits, losses = self(features)  # pylint: disable=not-callable
    if self._target_modality_is_real:
      return logits, logits, losses  # Raw numbers returned from real modality.
    if self.hparams.sampling_method == "argmax":
      samples = tf.argmax(logits, axis=-1)
    else:
      assert self.hparams.sampling_method == "random"

      def multinomial_squeeze(logits, temperature=1.0):
        logits_shape = common_layers.shape_list(logits)
        logits /= tf.reshape(temperature, [-1] + [1] * (len(logits_shape) - 1))
        reshaped_logits = tf.reshape(logits, [-1, logits_shape[-1]])
        choices = tf.multinomial(reshaped_logits, 1)
        choices = tf.reshape(choices, logits_shape[:-1])
        return choices

      temperature = features.get("sampling_temp", self.hparams.sampling_temp)
      samples = multinomial_squeeze(logits, temperature)

    return samples, logits, losses
//  [python]
def sample(self, features):
    """Run the model and extract samples.

    Args:
      features: an map of string to `Tensor`.

    Returns:
       samples: an integer `Tensor`.
       logits: a list of `Tensor`s, one per datashard.
       losses: a dictionary: {loss-name (string): floating point `Scalar`}.
    """
    logits, losses = self(features)  # pylint: disable=not-callable
    if self._target_modality_is_real:
      return logits, logits, losses  # Raw numbers returned from real modality.
    if self.hparams.sampling_method == "argmax":
      samples = tf.argmax(logits, axis=-1)
    else:
      assert self.hparams.sampling_method == "random"

      def multinomial_squeeze(logits, temperature=1.0):
        logits_shape = common_layers.shape_list(logits)
        logits /= tf.reshape(temperature, [-1] + [1] * (len(logits_shape) - 1))
        reshaped_logits = tf.reshape(logits, [-1, logits_shape[-1]])
        choices = tf.random.categorical(reshaped_logits, 1)
        choices = tf.reshape(choices, logits_shape[:-1])
        return choices

      temperature = features.get("sampling_temp", self.hparams.sampling_temp)
      samples = multinomial_squeeze(logits, temperature)

    return samples, logits, losses
//  [python]
def sample_with_temperature(logits, temperature):
  """Either argmax after softmax or random sample along the pitch axis.

  Args:
    logits: a Tensor of shape (batch, time, pitch, instrument).
    temperature: a float  0.0=argmax 1.0=random

  Returns:
    a Tensor of the same shape, with one_hots on the pitch dimension.
  """
  logits = tf.transpose(logits, [0, 1, 3, 2])
  pitch_range = tf.shape(logits)[-1]

  def sample_from_logits(logits):
    with tf.control_dependencies([tf.assert_greater(temperature, 0.0)]):
      logits = tf.identity(logits)
    reshaped_logits = (
        tf.reshape(logits, [-1, tf.shape(logits)[-1]]) / temperature)
    choices = tf.multinomial(reshaped_logits, 1)
    choices = tf.reshape(choices,
                         tf.shape(logits)[:logits.get_shape().ndims - 1])
    return choices

  choices = tf.cond(tf.equal(temperature, 0.0),
                    lambda: tf.argmax(tf.nn.softmax(logits), -1),
                    lambda: sample_from_logits(logits))
  samples_onehot = tf.one_hot(choices, pitch_range)
  return tf.transpose(samples_onehot, [0, 1, 3, 2])
//  [python]
def sample_with_temperature(logits, temperature):
  """Either argmax after softmax or random sample along the pitch axis.

  Args:
    logits: a Tensor of shape (batch, time, pitch, instrument).
    temperature: a float  0.0=argmax 1.0=random

  Returns:
    a Tensor of the same shape, with one_hots on the pitch dimension.
  """
  logits = tf.transpose(logits, [0, 1, 3, 2])
  pitch_range = tf.shape(logits)[-1]

  def sample_from_logits(logits):
    with tf.control_dependencies([tf.assert_greater(temperature, 0.0)]):
      logits = tf.identity(logits)
    reshaped_logits = (
        tf.reshape(logits, [-1, tf.shape(logits)[-1]]) / temperature)
    choices = tf.random.categorical(reshaped_logits, 1)
    choices = tf.reshape(choices,
                         tf.shape(logits)[:logits.get_shape().ndims - 1])
    return choices

  choices = tf.cond(tf.equal(temperature, 0.0),
                    lambda: tf.argmax(tf.nn.softmax(logits), -1),
                    lambda: sample_from_logits(logits))
  samples_onehot = tf.one_hot(choices, pitch_range)
  return tf.transpose(samples_onehot, [0, 1, 3, 2])
//  [python]
def generate(self, x=None, y=None, h=None, mode='random'):
    '''
    Output 1 generated sample point.

    Input:
      `x`: wav signal, `int`, [b, T]
      `y`: phone sequence (local conditional signal), `int`, [b, T]
      `h`: speaker id (global conditional signal), `int`, [b,]
      `mode`:
        `random` (default; output randomly from 256 possibilities proportional to logits)
        `greedy`

    Output:
      `xh`: generated sample of next time slot, `int`, [b, 1]

    FIXME: Generation is insanely slow. It could be made much more efficient
           if we don't re-compute the already available values.
           Currently, it takes 10 minutes to generate [4, 16000] samples.
    '''
    print('Decoder output mode: {}'.format(mode))
    with tf.name_scope('IncrementalGeneration'):
      x_analog = self._D2A(x)
      x_next = self._Dec(x_analog, y, h)
      x_next = x_next[:, -1, :]

      # The output mode
      if mode == 'random':
        return tf.multinomial(x_next, 1)  # [b, 1], int64
      elif mode == 'greedy':
        return tf.argmax(x_next, axis=-1)
      else:
        raise NotImplementedError('"Mode: {}" is not implemented'.format(mode))
//  [python]
def generate(self, x=None, y=None, h=None, mode='random'):
    '''
    Output 1 generated sample point.

    Input:
      `x`: wav signal, `int`, [b, T]
      `y`: phone sequence (local conditional signal), `int`, [b, T]
      `h`: speaker id (global conditional signal), `int`, [b,]
      `mode`:
        `random` (default; output randomly from 256 possibilities proportional to logits)
        `greedy`

    Output:
      `xh`: generated sample of next time slot, `int`, [b, 1]

    FIXME: Generation is insanely slow. It could be made much more efficient
           if we don't re-compute the already available values.
           Currently, it takes 10 minutes to generate [4, 16000] samples.
    '''
    print('Decoder output mode: {}'.format(mode))
    with tf.name_scope('IncrementalGeneration'):
      x_analog = self._D2A(x)
      x_next = self._Dec(x_analog, y, h)
      x_next = x_next[:, -1, :]

      # The output mode
      if mode == 'random':
        return tf.random.categorical(x_next, 1)  # [b, 1], int64
      elif mode == 'greedy':
        return tf.argmax(x_next, axis=-1)
      else:
        raise NotImplementedError('"Mode: {}" is not implemented'.format(mode))
//  [python]
def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0, top_p=0.0):
    if start_token is None:
        assert context is not None, 'Specify exactly one of start_token and context!'
    else:
        assert context is None, 'Specify exactly one of start_token and context!'
        context = tf.fill([batch_size, 1], start_token)

    def step(hparams, tokens, past=None):
        lm_output = model.model(hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)

        logits = lm_output['logits'][:, :, :hparams.n_vocab]
        presents = lm_output['present']
        presents.set_shape(model.past_shape(hparams=hparams, batch_size=batch_size))
        return {
            'logits': logits,
            'presents': presents,
        }

    with tf.name_scope('sample_sequence'):
        # Don't feed the last context token -- leave that to the loop below
        # TODO: Would be slightly faster if we called step on the entire context,
        # rather than leaving the last token transformer calculation to the while loop.
        context_output = step(hparams, context[:, :-1])

        def body(past, prev, output):
            next_outputs = step(hparams, prev[:, tf.newaxis], past=past)
            logits = next_outputs['logits'][:, -1, :]  / tf.to_float(temperature)
            if top_p > 0.0:
                logits = top_p_logits(logits, p=top_p)
            else:
                logits = top_k_logits(logits, k=top_k)
            samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32)
            return [
                tf.concat([past, next_outputs['presents']], axis=-2),
                tf.squeeze(samples, axis=[1]),
                tf.concat([output, samples], axis=1),
            ]

        def cond(*args):
            return True

        _, _, tokens = tf.while_loop(
            cond=cond, body=body,
            maximum_iterations=length,
            loop_vars=[
                context_output['presents'],
                context[:, -1],
                context,
            ],
            shape_invariants=[
                tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
                tf.TensorShape([batch_size]),
                tf.TensorShape([batch_size, None]),
            ],
            back_prop=False,
        )

        return tokens
//  [python]
def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0, top_p=0.0):
    if start_token is None:
        assert context is not None, 'Specify exactly one of start_token and context!'
    else:
        assert context is None, 'Specify exactly one of start_token and context!'
        context = tf.fill([batch_size, 1], start_token)

    def step(hparams, tokens, past=None):
        lm_output = model.model(hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)

        logits = lm_output['logits'][:, :, :hparams.n_vocab]
        presents = lm_output['present']
        presents.set_shape(model.past_shape(hparams=hparams, batch_size=batch_size))
        return {
            'logits': logits,
            'presents': presents,
        }

    with tf.name_scope('sample_sequence'):
        # Don't feed the last context token -- leave that to the loop below
        # TODO: Would be slightly faster if we called step on the entire context,
        # rather than leaving the last token transformer calculation to the while loop.
        context_output = step(hparams, context[:, :-1])

        def body(past, prev, output):
            next_outputs = step(hparams, prev[:, tf.newaxis], past=past)
            logits = next_outputs['logits'][:, -1, :]  / tf.to_float(temperature)
            if top_p > 0.0:
                logits = top_p_logits(logits, p=top_p)
            else:
                logits = top_k_logits(logits, k=top_k)
            samples = tf.random.categorical(logits, num_samples=1, dtype=tf.int32)
            return [
                tf.concat([past, next_outputs['presents']], axis=-2),
                tf.squeeze(samples, axis=[1]),
                tf.concat([output, samples], axis=1),
            ]

        def cond(*args):
            return True

        _, _, tokens = tf.while_loop(
            cond=cond, body=body,
            maximum_iterations=length,
            loop_vars=[
                context_output['presents'],
                context[:, -1],
                context,
            ],
            shape_invariants=[
                tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
                tf.TensorShape([batch_size]),
                tf.TensorShape([batch_size, None]),
            ],
            back_prop=False,
        )

        return tokens
//  [python]
def _preprocess(self, features):
    """Preprocesses features for multilingual translation."""
    seqs, tags = {}, {}

    if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
      seqs["src"] = features["inputs"]
      seqs["tgt"] = features["targets"]
      seqs["aux"] = None
      tags["src"] = features["input_tags"]
      tags["tgt"] = features["target_tags"]
      tags["aux"] = None

      # Construct a tensor of auxiliary tags.
      batch_size = common_layers.shape_list(features["all_tags"])[0]
      num_all_tags = common_layers.shape_list(features["all_tags"])[1]
      # <float32> [num_all_tags, 1, emb_dim].
      all_tags = features["all_tags"][0]  # batch elements are identical.
      # <int32> [batch_size].
      aux_tag_index = tf.multinomial(
          tf.ones([1, num_all_tags]), batch_size,
          output_dtype=tf.int32)[0]
      # <float32> [batch_size, 1, 1, emb_dim].
      tags["aux"] = tf.expand_dims(tf.gather(all_tags, aux_tag_index), 1)

      from_domains = ["src", "src", "tgt"]
      to_domains = ["tgt", "aux", "aux"]
    else:
      seqs["src"] = features["inputs"]
      seqs["tgt"] = features["targets"]
      tags["src"] = None
      tags["tgt"] = features["target_tags"]

      # Expand target tags to beam width, if necessary.
      if self._hparams.mode == tf_estimator.ModeKeys.PREDICT:
        tags["tgt"] = tf.tile(tags["tgt"], [self._hparams.beam_width, 1, 1, 1])

      from_domains = ["src"]
      to_domains = ["tgt"]

    # Construct inputs and targets.
    inputs, targets = {}, {}
    for fd, td in zip(from_domains, to_domains):
      key = "%s>%s" % (fd, td)
      inputs[key], targets[key] = self._build_inputs_and_targets(
          seqs[fd], tags[fd], seqs[td], tags[td])

    return inputs, targets
//  [python]
def _preprocess(self, features):
    """Preprocesses features for multilingual translation."""
    seqs, tags = {}, {}

    if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
      seqs["src"] = features["inputs"]
      seqs["tgt"] = features["targets"]
      seqs["aux"] = None
      tags["src"] = features["input_tags"]
      tags["tgt"] = features["target_tags"]
      tags["aux"] = None

      # Construct a tensor of auxiliary tags.
      batch_size = common_layers.shape_list(features["all_tags"])[0]
      num_all_tags = common_layers.shape_list(features["all_tags"])[1]
      # <float32> [num_all_tags, 1, emb_dim].
      all_tags = features["all_tags"][0]  # batch elements are identical.
      # <int32> [batch_size].
      aux_tag_index = tf.random.categorical(
          tf.ones([1, num_all_tags]), batch_size,
          dtype=tf.int32)[0]
      # <float32> [batch_size, 1, 1, emb_dim].
      tags["aux"] = tf.expand_dims(tf.gather(all_tags, aux_tag_index), 1)

      from_domains = ["src", "src", "tgt"]
      to_domains = ["tgt", "aux", "aux"]
    else:
      seqs["src"] = features["inputs"]
      seqs["tgt"] = features["targets"]
      tags["src"] = None
      tags["tgt"] = features["target_tags"]

      # Expand target tags to beam width, if necessary.
      if self._hparams.mode == tf_estimator.ModeKeys.PREDICT:
        tags["tgt"] = tf.tile(tags["tgt"], [self._hparams.beam_width, 1, 1, 1])

      from_domains = ["src"]
      to_domains = ["tgt"]

    # Construct inputs and targets.
    inputs, targets = {}, {}
    for fd, td in zip(from_domains, to_domains):
      key = "%s>%s" % (fd, td)
      inputs[key], targets[key] = self._build_inputs_and_targets(
          seqs[fd], tags[fd], seqs[td], tags[td])

    return inputs, targets
//  [python]
def _cat_sample(self, sample_shape=(), seed=None, name='sample'):
  assert len(sample_shape) in (0, 1), sample_shape
  assert len(self.logits_parameter().shape) == 2
  indices = tf.multinomial(
      self.logits_parameter(), sample_shape[0] if sample_shape else 1,
      seed, name)
  if not sample_shape:
    indices = indices[..., 0]
  return indices
//  [python]
def _cat_sample(self, sample_shape=(), seed=None, name='sample'):
  assert len(sample_shape) in (0, 1), sample_shape
  assert len(self.logits_parameter().shape) == 2
  indices = tf.random.categorical(
      self.logits_parameter(), sample_shape[0] if sample_shape else 1,
      self.dtype, seed, name)
  if not sample_shape:
    indices = indices[..., 0]
  return indices
//  [python]
def _cat_sample(self, sample_shape=(), seed=None, name="sample"):
    assert len(sample_shape) in (0, 1), sample_shape
    assert len(self.logits_parameter().shape) == 2
    indices = tf.multinomial(
        self.logits_parameter(), sample_shape[0] if sample_shape else 1, seed, name
    )
    if not sample_shape:
        indices = indices[..., 0]
    return indices